From 5733bbc834d5df9dc727dae994694144f0aa4a50 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:39 +0000 Subject: [PATCH] refactor(redis): move Claude session router bindings behind a typed store --- .../stores/claude_code_router_binding.py | 69 + litellm/router.py | 93 +- .../router_code_coverage.py | 2 + .../stores/test_claude_code_router_binding.py | 223 ++ tests/unit/test_router/test_router.py | 2503 +++++++++-------- 5 files changed, 1633 insertions(+), 1257 deletions(-) create mode 100644 litellm/redis/stores/claude_code_router_binding.py create mode 100644 tests/unit/redis/stores/test_claude_code_router_binding.py diff --git a/litellm/redis/stores/claude_code_router_binding.py b/litellm/redis/stores/claude_code_router_binding.py new file mode 100644 index 00000000000..048c01f18ae --- /dev/null +++ b/litellm/redis/stores/claude_code_router_binding.py @@ -0,0 +1,69 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_PREFIX: Final = "claude_code_session_router:v1:" +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_OBJECT: Final = TypeAdapter(object) + + +@dataclass(frozen=True, slots=True) +class SessionRouterScope: + caller: str + session_id: str + + +def binding_key(scope: SessionRouterScope) -> str: + return f"{_PREFIX}{scope.caller}:{scope.session_id}" + + +def binding_scope(key: str) -> SessionRouterScope | None: + if not key.startswith(_PREFIX): + return None + caller, separator, session_id = key[len(_PREFIX) :].rpartition(":") + return SessionRouterScope(caller, session_id) if separator and caller and session_id else None + + +def _decode(raw: bytes) -> object: + text: Final = raw.decode("utf-8") + try: + return _JSON.validate_json(text) + except ValueError: + return _OBJECT.validate_python(ast.literal_eval(text)) + + +class ClaudeCodeRouterBindingStore(Protocol): + async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure: ... + + async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure: ... + + async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportClaudeCodeRouterBindingStore: + transport: RedisTransport + + async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure: + raw: Final = await self.transport.get(binding_key(scope)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + value: Final = _decode(raw) + except (ValueError, SyntaxError, UnicodeError) as error: + return RedisFailure(kind="decode", message=str(error)) + return value if isinstance(value, str) else None + + async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure: + result: Final = await self.transport.set_seconds(binding_key(scope), json.dumps(model), ttl=ttl) + return result if isinstance(result, RedisFailure) else None + + async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure: + result: Final = await self.transport.delete(binding_key(scope)) + return result if isinstance(result, RedisFailure) else None diff --git a/litellm/router.py b/litellm/router.py index 1809f861ad8..76275bb9618 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -83,7 +83,6 @@ from litellm.caching.caching import ( RedisCache, RedisClusterCache, ) -from litellm.caching.redis_cache import log_redis_failure from litellm.constants import ( BAGGAGE_MAPPED_TAGS_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, @@ -182,6 +181,13 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.redis.stores.claude_code_router_binding import ( + ClaudeCodeRouterBindingStore, + SessionRouterScope, + TransportClaudeCodeRouterBindingStore, + binding_key, + binding_scope, +) from litellm.redis.stores.router_global_usage import ( GlobalUsageCounter, RouterGlobalUsageStore, @@ -16238,6 +16244,10 @@ class Router(DeprecatedRouterMembers): ) def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: + scope: Final = self._claude_code_session_router_scope(request_kwargs) + return None if scope is None else binding_key(scope) + + def _claude_code_session_router_scope(self, request_kwargs: Mapping[str, object]) -> SessionRouterScope | None: session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: return None @@ -16248,33 +16258,52 @@ class Router(DeprecatedRouterMembers): caller_scope: Final = metadata.get("user_api_key_hash") if not isinstance(caller_scope, str) or not caller_scope: return None - return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + return SessionRouterScope(caller_scope, session_id) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + async def _delete_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> None: + scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key + if scope is None: + return try: - await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", - e, - ) + session_cache: Final = self._claude_code_session_router_cache + session_cache.in_memory_cache.delete_cache(binding_key(scope)) + redis: Final = session_cache.redis_cache + if redis is not None: + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + await store.remove_binding(scope) + except Exception as e: # noqa: BLE001 # local cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning("Failed to delete Claude Code session router binding: %s", e) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: - session_cache: Final = self._claude_code_session_router_cache - try: - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis - log_redis_failure( - verbose_router_logger, - logging.WARNING, - "Failed to read Claude Code session router binding; using the requested model", - e, - ) + async def _get_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> str | None: + scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key + if scope is None: return None + session_cache: Final = self._claude_code_session_router_cache + redis: Final = session_cache.redis_cache + if redis is None: + local: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + await session_cache.async_get_cache(key=binding_key(scope), local_only=True) + ) + return local if isinstance(local, str) else None + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + shared: Final = await store.read_binding(scope) + return None if isinstance(shared, RedisFailure) else shared + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _bind_claude_code_session_router(self, scope: SessionRouterScope, model: str) -> None: + session_cache: Final = self._claude_code_session_router_cache + await session_cache.async_set_cache( + key=binding_key(scope), + value=model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + local_only=True, + ) + redis: Final = session_cache.redis_cache + if redis is not None: + store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport()) + await store.bind_router(scope, model, ttl=timedelta(seconds=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS)) @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) async def _resolve_claude_code_session_router( @@ -16294,36 +16323,28 @@ class Router(DeprecatedRouterMembers): ) ): return registered_model_name - cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) - if cache_key is None or not isinstance(request_kwargs, dict): + scope: Final = self._claude_code_session_router_scope(request_kwargs) + if scope is None or not isinstance(request_kwargs, dict): return registered_model_name if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) + bound_model: Final = await self._get_claude_code_session_router_binding(scope) if not isinstance(bound_model, str): return registered_model_name bound_registered_model: Final = self._core_get_model_from_alias(model=bound_model) or bound_model if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._delete_claude_code_session_router_binding(scope) return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=bound_model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) + await self._bind_claude_code_session_router(scope, bound_model) self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) return bound_registered_model if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) + await self._bind_claude_code_session_router(scope, model) return registered_model_name async def _core_async_pre_routing_hook( diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 791515f3b41..9a5a37d53ef 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -177,6 +177,8 @@ ignored_function_names = [ "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_read_global_usage_remote", # Tested through the shared Redis usage header tests in test_router_helper_utils.py "_set_global_usage_local", # Tested through the shared Redis usage header tests in test_router_helper_utils.py + "_claude_code_session_router_scope", # Tested through the Claude Code session router binding tests in test_router.py + "_bind_claude_code_session_router", # Tested through the Claude Code session router binding tests in test_router.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py "_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py diff --git a/tests/unit/redis/stores/test_claude_code_router_binding.py b/tests/unit/redis/stores/test_claude_code_router_binding.py new file mode 100644 index 00000000000..35c45c74b5f --- /dev/null +++ b/tests/unit/redis/stores/test_claude_code_router_binding.py @@ -0,0 +1,223 @@ +import json +from collections.abc import Iterator, Mapping +from datetime import timedelta +from typing import Final, Protocol, cast +from unittest.mock import patch + +import fakeredis +import pytest +from pydantic import TypeAdapter +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm import Router +from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.claude_code_router_binding import SessionRouterScope, TransportClaudeCodeRouterBindingStore +from litellm.redis.transport.contract import Arg, RedisFailure +from litellm.types.router import PreRoutingHookResponse +from tests._support.fake_redis import fake_redis_cache + +KEY: Final = "claude_code_session_router:v1:key:scope:session-1234" +MODEL: Final = 'smart-雪 "router"' +SCOPE: Final = SessionRouterScope("key:scope", "session-1234") +_METADATA: Final = TypeAdapter(dict[str, object]) + + +class BindingRouter(Protocol): + def update_redis_cache(self, cache: RedisCache) -> None: ... + + async def async_pre_routing_hook( + self, model: str, request_kwargs: Mapping[str, object] + ) -> PreRoutingHookResponse | None: ... + + +class LegacyBindingCache(Protocol): + async def async_get_cache(self, key: str) -> object: ... + + async def async_set_cache(self, key: str, value: str, *, ttl: int) -> None: ... + + +@pytest.fixture(autouse=True) +def clock() -> Iterator[None]: + with patch("time.time", return_value=1800000000.0): + yield + + +def _router(cache: RedisCache | None) -> BindingRouter: + router: Final = Router( + model_list=[ + {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}}, + *( + { + "model_name": model, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "plain") + }, + }, + } + for model in (MODEL, "premium-router") + ), + ] + ) + typed: Final = cast(BindingRouter, router) # cast-ok: typed view of the actual binding caller, without replacing it + if cache is not None: + typed.update_redis_cache(cache) + return typed + + +def _kwargs( + *, agent: bool = False, caller: str = "key:scope", session: str = "session-1234", depth: int = 0 +) -> Mapping[str, object]: + return { + "metadata": {"user_api_key_hash": caller}, + "proxy_server_request": { + "headers": { + "X-Claude-Code-Session-Id": session, + **({"x-claude-code-agent-id": "agent-1234"} if agent else {}), + } + }, + "fallback_depth": depth, + } + + +async def _route(router: BindingRouter, model: str, kwargs: Mapping[str, object]) -> str: + await router.async_pre_routing_hook(model, dict(kwargs)) + bound: Final = _METADATA.validate_python(kwargs["metadata"]).get("model_group") + return bound if isinstance(bound, str) else model + + +@pytest.mark.asyncio +async def test_binding_old_new_wire_namespace_and_read_refresh(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + legacy: Final = cast(LegacyBindingCache, cache) # cast-ok: rolling-deploy writer and reader + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await legacy.async_set_cache(KEY, MODEL, ttl=3600) + router: Final = _router(cache) + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + assert await raw.get("tenant:" + KEY) == json.dumps(MODEL).encode() + assert await raw.exists(KEY) == 0 + with patch("time.time", return_value=1800000100.0): + assert await raw.pttl("tenant:" + KEY) == 3500000 + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + assert await raw.pttl("tenant:" + KEY) == 3600000 + assert await _route(router, "premium-router", _kwargs()) == "premium-router" + assert await legacy.async_get_cache(KEY) == "premium-router" + assert await raw.get("tenant:" + KEY) == b'"premium-router"' + + +@pytest.mark.asyncio +async def test_shared_miss_and_remote_update_override_a_fresh_local_binding(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + main: Final = _router(cache) + agent: Final = _router(cache) + await _route(main, MODEL, _kwargs()) + assert await _route(agent, "plain", _kwargs(agent=True)) == MODEL + await _route(main, "premium-router", _kwargs()) + assert await _route(agent, "plain", _kwargs(agent=True)) == "premium-router" + await raw.delete(KEY) + assert await _route(agent, "plain", _kwargs(agent=True)) == "plain" + + +@pytest.mark.asyncio +async def test_local_bindings_survive_without_redis_but_are_not_backfilled_on_attach( + server: fakeredis.FakeServer, +) -> None: + router: Final = _router(None) + await _route(router, MODEL, _kwargs()) + assert await _route(router, "plain", _kwargs(agent=True)) == MODEL + router.update_redis_cache(fake_redis_cache(server)) + assert await _route(router, "plain", _kwargs(agent=True)) == "plain" + assert await fakeredis.FakeAsyncRedis(server=server).exists(KEY) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_value", [b"null", b"4", b"[]", b'{"router": "other"}', b"not-json", b"\xff"]) +async def test_non_string_and_malformed_bindings_are_ignored_not_rewritten( + server: fakeredis.FakeServer, raw_value: bytes +) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, raw_value, ex=600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain" + assert await raw.get(KEY) == raw_value + assert await raw.pttl(KEY) == 600000 + + +@pytest.mark.asyncio +async def test_legacy_python_literal_binding_refreshes_as_json(server: fakeredis.FakeServer) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, b"'premium-router'", ex=600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "premium-router" + assert await raw.get(KEY) == b'"premium-router"' + assert await raw.pttl(KEY) == 3600000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [b'"removed-router"', b'""']) +async def test_unknown_strategy_deletes_the_binding_instead_of_tombstoning( + server: fakeredis.FakeServer, value: bytes +) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(KEY, value, ex=3600) + assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain" + assert await raw.exists(KEY) == 0 + + +@pytest.mark.asyncio +async def test_invalid_scope_other_caller_and_fallback_cannot_reuse_the_binding(server: fakeredis.FakeServer) -> None: + router: Final = _router(fake_redis_cache(server)) + await _route(router, MODEL, _kwargs()) + assert await _route(router, "plain", _kwargs(agent=True, caller="other")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, caller="")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, session="short")) == "plain" + assert await _route(router, "plain", _kwargs(agent=True, depth=1)) == "plain" + assert await fakeredis.FakeAsyncRedis(server=server).get(KEY) == json.dumps(MODEL).encode() + + +@pytest.mark.asyncio +async def test_binding_outage_never_uses_stale_local_and_shares_the_breaker() -> None: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + router: Final = _router(cache) + await _route(router, MODEL, _kwargs()) + server.connected = False + assert await _route(router, "plain", _kwargs(agent=True)) == "plain" + assert await _route(router, "premium-router", _kwargs()) == "premium-router" + store: Final = TransportClaudeCodeRouterBindingStore(cache.transport()) + for _ in range(5): + await store.read_binding(SCOPE) + result: Final = await store.bind_router(SCOPE, MODEL, ttl=timedelta(seconds=3600)) + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +class RedisWriter(Protocol): + async def set( + self, name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False + ) -> object: ... + + +@pytest.mark.asyncio +async def test_unknown_applied_binding_write_remains_routable_without_replay(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + router: Final = _router(cache) + original_set: Final = cast(type[RedisWriter], fakeredis.FakeAsyncRedis).set # cast-ok: typed redis-py boundary + + async def apply_then_disconnect( + client: RedisWriter, + name: str, + value: Arg, + *, + ex: timedelta | None = None, + px: int | None = None, + nx: bool = False, + ) -> object: + await original_set(client, name, value, ex=ex, px=px, nx=nx) + raise RedisConnectionError("response lost after SET applied") + + with patch.object(fakeredis.FakeAsyncRedis, "set", new=apply_then_disconnect): + assert await _route(router, MODEL, _kwargs()) == MODEL + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(KEY) == json.dumps(MODEL).encode() + assert await raw.pttl(KEY) == 3600000 + assert await _route(_router(cache), "plain", _kwargs(agent=True)) == MODEL diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a6436e47507..3228feff84f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1,16 +1,21 @@ -import ast, asyncio, importlib, time +import ast +import asyncio import contextlib import copy import functools import gc -import itertools +import importlib +import inspect import io +import itertools import json import linecache import logging import os import sys import threading +import time +import traceback import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timedelta, timezone @@ -25,45 +30,48 @@ import pytest import respx from fastapi import HTTPException from opentelemetry import trace -from pydantic import BaseModel, TypeAdapter +from pydantic import TypeAdapter import litellm from litellm import APIConnectionError, Router from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching import RedisCache, RedisClusterCache -from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache -from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.plumbing.context import set_request_root_span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) +from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) from litellm.llms.bedrock.common_utils import BedrockError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth +from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, - _stream_commits_now, + Span, _anthropic_stream_error_is_gateway_verdict, - _stream_fallback_error_for_raised, _anthropic_stream_forwards_ping_live, - _stream_raised_error_status, _anthropic_stream_should_decline_fallback, _is_retriable_anthropic_status, _responses_stream_holds_event, - _without_line_breaks, _search_tool_support, - Span, + _stream_commits_now, + _stream_fallback_error_for_raised, + _stream_raised_error_status, + _without_line_breaks, ) -from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.savings_baseline import Baseline, _models_in @@ -94,19 +102,11 @@ from litellm.types.router import ( RouterRateLimitError, RoutingContext, ) -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.vector_store.transformation import( - LiteLLMVectorStoreEmbeddingExecutor, - RouterVectorStoreEmbeddingExecutor, -) from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from tests.large_text import text -import traceback -import inspect -from typing import List, Optional if TYPE_CHECKING: from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -351,31 +351,18 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -406,13 +393,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -453,10 +436,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -476,9 +456,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -525,18 +503,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -547,9 +517,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -759,7 +727,9 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards from io import BytesIO from unittest.mock import MagicMock, patch - jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + jsonl_content = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + ) router = litellm.Router( model_list=[ { @@ -770,9 +740,7 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards ) with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: - await router.acreate_file( - model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True - ) + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True) forwarded = mock_acreate_file.call_args.kwargs assert forwarded["passthrough"] is True forwarded["file"].seek(0) @@ -1044,7 +1012,8 @@ async def test_async_router_afile_content_uses_deployment_custom_llm_provider(): This prevents "None is not a valid LlmProviders" errors when calling file content operations. """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import MagicMock, patch + from litellm.types.llms.openai import HttpxBinaryResponseContent router = litellm.Router( @@ -1248,9 +1217,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -1510,6 +1477,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1521,8 +1489,7 @@ async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: moved = sorted( f"{key}={cache_dict[key]}" for key in cache_dict - if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) - and cache_dict[key] + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) and cache_dict[key] ) if moved: return moved @@ -1603,9 +1570,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1666,7 +1631,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1750,9 +1715,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1760,9 +1723,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1778,30 +1741,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1810,9 +1761,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1863,9 +1812,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1939,7 +1886,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1964,7 +1911,9 @@ def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" ) - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}] + ) monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) return "test-onprem-llm" @@ -2041,15 +1990,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -2084,7 +2027,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -2095,7 +2038,6 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): assert "No deployment available" in str(exc_info.value) # Test 4: Test with semaphore (rate limiting) - import asyncio async def mock_semaphore_function(**kwargs): return {"result": "semaphore_success", "model": kwargs.get("model")} @@ -2114,15 +2056,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo" ) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -2151,16 +2087,10 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with pytest.raises(Exception, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -2234,9 +2164,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) @pytest.mark.asyncio @@ -2342,14 +2272,10 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] @@ -2444,9 +2370,7 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -2474,9 +2398,7 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -2962,9 +2884,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -3010,9 +2930,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "_core_get_deployment_model_info", side_effect=Exception("not found")): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -3272,6 +3190,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -3335,7 +3254,7 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi """Flag on + plain-text break: the router re-enters the fallback chain with an assistant-prefill continuation and the mid-stream marker, then streams the fallback's output instead of re-raising.""" - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, patch from litellm.exceptions import MidStreamFallbackError from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( @@ -3404,8 +3323,12 @@ async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_ enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="123-45-", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) @@ -3450,8 +3373,12 @@ async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_rec def cut(generated_content): return MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=generated_content, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) messages = [{"role": "user", "content": "Hi"}] @@ -3461,7 +3388,11 @@ async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_rec first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), ), messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], - initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER, + }, ) backup_leg = _FakeFallbackStream( litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) @@ -3500,8 +3431,12 @@ async def test_acompletion_streaming_iterator_continuation_never_repeats_the_par enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) async def fallback_stream(): @@ -3537,8 +3472,12 @@ async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_ enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) primary_chunk = litellm.ModelResponseStream( id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] @@ -3623,8 +3562,12 @@ async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs( messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] partial = "Transformers use " error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=partial, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) content_chunk = litellm.ModelResponseStream( model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] @@ -3646,20 +3589,35 @@ class _CutSSE(httpx.AsyncByteStream): self._content: Final = content async def __aiter__(self) -> AsyncIterator[bytes]: - yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + yield f"data: {json.dumps({'id': 'cut', 'object': 'chat.completion.chunk', 'created': 1700000000, 'model': 'm', 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': self._content}, 'finish_reason': None}]})}\n\n".encode() raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") -_BACKUP_SSE: Final = "".join( - f"data: {json.dumps(chunk)}\n\n" - for chunk in ( - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]}, - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}}, +_BACKUP_SSE: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None} + ], + }, + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}, + }, + ) ) -) + "data: [DONE]\n\n" + + "data: [DONE]\n\n" +) class _SpendRecorder(CustomLogger): @@ -3701,7 +3659,9 @@ async def test_continued_stream_success_log_bills_every_cut_leg( recorder: Final = _SpendRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.deepseek.com").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) @@ -3784,7 +3744,9 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t recorder: Final = _SlowFailureRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.mistral.ai").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) @@ -3899,7 +3861,10 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall 200, headers={"content-type": "text/event-stream"}, text=_responses_sse( - {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.created", + "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress"), + }, { "type": "response.completed", "response": _responses_envelope( @@ -3973,8 +3938,13 @@ def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherit from litellm.types.utils import Usage cut_logging: Final = LiteLLMLogging( - model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="cut", function_id="cut", + model="deepseek/deepseek-chat", + messages=[], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="cut", + function_id="cut", ) cut_logging.model_call_details["response_cost"] = 0.5 cut_logging.update_environment_variables( @@ -3986,8 +3956,13 @@ def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherit response_id="chatcmpl-primary", ) fallback_logging: Final = LiteLLMLogging( - model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", + model="mistral-large-latest", + messages=[], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="fallback", + function_id="fallback", ) Router._carry_partial_spend_to_fallback( @@ -4016,7 +3991,10 @@ async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_ from litellm.exceptions import MidStreamFallbackError error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", + message="Connection lost", + model="gpt-4", + llm_provider="anthropic", + generated_content="", is_pre_first_chunk=True, ) role_chunk = litellm.ModelResponseStream( @@ -4050,8 +4028,12 @@ async def test_acompletion_streaming_iterator_continuation_does_not_double_trail from litellm.exceptions import MidStreamFallbackError error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content=partial, + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_items = [ litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts @@ -4088,8 +4070,12 @@ async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, ena router.update_settings(enable_mid_stream_fallback_continuation=enabled) error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) with patch.object( @@ -4137,8 +4123,12 @@ async def test_acompletion_streaming_iterator_declines_ineligible_after_content( enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, **error_kwargs, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + **error_kwargs, ) source = _make_midstream_source(error) initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} @@ -4172,8 +4162,13 @@ async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_c ) content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") error = MidStreamFallbackError( - message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", - is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, + message="filtered", + model="gpt-4", + llm_provider="azure", + generated_content="Step one is to obtain the", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, + original_exception=content_filter, ) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: @@ -4232,8 +4227,12 @@ async def _continues_after_content( enable_mid_stream_fallback_continuation=True, ) error: Final = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) cut_stream: Final = _make_midstream_source(error) @@ -4448,14 +4447,19 @@ async def test_acompletion_streaming_iterator_flag_off_declines_after_content(fl **flag_kwargs, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) source = _make_midstream_source(error) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], + model_response=source, + messages=[{"role": "user", "content": "Hi"}], initial_kwargs={"model": "gpt-4", "stream": True}, ) with pytest.raises(MidStreamFallbackError): @@ -4545,8 +4549,12 @@ async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plai enable_mid_stream_fallback_continuation=True, ) error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: @@ -4577,7 +4585,10 @@ def test_build_completion_continuation_input_keeps_structured_assistant_content( assert [m["role"] for m in built] == ["user", "assistant"] assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} - assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] + assert _extend_trailing_text_block([image_block], "here it is") == [ + image_block, + {"type": "text", "text": "here it is"}, + ] assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] assert ( litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") @@ -4595,11 +4606,18 @@ def test_mid_stream_continuation_eligible_allows_text_response_format(): enable_mid_stream_fallback_continuation=True, ) e = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, + message="boom", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, ) assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False + assert ( + router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) + is False + ) def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): @@ -4607,8 +4625,11 @@ def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") wrapped = MidStreamFallbackError( - message="rate limited", model="gpt-4", llm_provider="openai", - original_exception=provider_error, generated_content="Hello", + message="rate limited", + model="gpt-4", + llm_provider="openai", + original_exception=provider_error, + generated_content="Hello", ) with pytest.raises(RateLimitError) as raised: litellm.Router._raise_original_mid_stream_error(wrapped) @@ -4955,9 +4976,7 @@ def test_adopt_fallback_response_headers_replaces_rather_than_merges(): "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, } - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -5029,9 +5048,7 @@ def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_repla fallback._response_headers = None fallback._hidden_params = {"model_id": "fallback-deployment"} - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params["model_id"] == "fallback-deployment" @@ -5058,9 +5075,7 @@ def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none() hidden_params_before = wrapper._hidden_params fallback = object() - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) + wrapper.adopt_fallback_response_headers(fallback, Router._prepare_fallback_hidden_params(fallback)) assert wrapper._response_headers is None assert wrapper._hidden_params is hidden_params_before @@ -5134,9 +5149,7 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro **additional_headers, ) - assert not { - key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") - } + assert not {key for key in proxy_headers if key.startswith("x-litellm-complexity-router-")} @pytest.mark.asyncio @@ -6312,11 +6325,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -6416,9 +6425,7 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) + router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -6893,9 +6900,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + assert "model_group" not in fbk.get("metadata", {}), ( + "model_group leaked into 'metadata' instead of 'litellm_metadata'" + ) @pytest.mark.asyncio @@ -7462,7 +7469,6 @@ async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_ assert backup_route.call_count == 0 - _BACKGROUND_TEXT_DELTA: Final = { "type": "response.output_text.delta", "sequence_number": 2, @@ -7533,6 +7539,7 @@ async def test_aresponses_background_stream_content_policy_error_reads_the_same_ assert primary_route.call_count == 2 assert backup_route.call_count == 0 + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" @@ -7610,9 +7617,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) + fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -7621,9 +7626,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), + return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), ), patch.object( router, @@ -7924,9 +7927,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -7954,14 +7955,10 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8068,14 +8065,10 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8102,9 +8095,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -8200,14 +8191,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8235,9 +8222,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counting_threads = [] @@ -8272,9 +8257,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8310,9 +8293,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) count = await router._acount_pre_call_check_tokens( model="m", @@ -8340,9 +8321,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8367,9 +8346,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -8411,9 +8388,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -8524,7 +8499,9 @@ def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): messages = [{"role": "user", "content": "hi"}] messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) + monkeypatch.setattr( + router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens} + ) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -8604,7 +8581,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -8642,9 +8619,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -8735,9 +8710,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -8745,15 +8718,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -8764,26 +8733,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -8802,9 +8763,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -8824,9 +8783,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -8859,9 +8816,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -8870,12 +8825,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -8898,9 +8849,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -8936,15 +8885,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -8991,18 +8936,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -9012,29 +8953,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -9051,7 +8980,11 @@ def test_get_deployment_model_info_base_model_merge_priority(): [ ( "gpt", - {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, + { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.services.ai.azure.com", + "api_key": "key", + }, "gpt/openai/deployments/gpt-5.4-mini/chat/completions", "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", ), @@ -9106,10 +9039,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( + f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + ) # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -9121,10 +9053,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( + f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + ) # Test Case 3: Bedrock converse endpoint kwargs = { @@ -9136,9 +9067,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + ) # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -9149,9 +9080,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + ) def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -9288,14 +9219,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -9460,7 +9387,9 @@ def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), + pytest.param( + True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern" + ), pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), ], @@ -9469,7 +9398,13 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( with_catch_all: bool, requested: str, team_id: str | None, expected_ids: tuple[str, ...] ): catch_all: Final = ( - ({"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, "model_info": {"id": "catch-all"}},) + ( + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, + "model_info": {"id": "catch-all"}, + }, + ) if with_catch_all else () ) @@ -9490,7 +9425,9 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( ) served: Final = router.wildcard_route_deployments(requested, team_id) - early: Final = router._try_early_resolve_deployments_for_model_not_in_names(model=requested, request_team_id=team_id) + early: Final = router._try_early_resolve_deployments_for_model_not_in_names( + model=requested, request_team_id=team_id + ) early_deployments: Final = early[1] if early is not None and isinstance(early[1], list) else [] assert tuple(Deployment.model_validate(d).model_info.id for d in served) == expected_ids @@ -9553,15 +9490,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -9588,9 +9520,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -9675,9 +9605,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -9714,9 +9642,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -9760,9 +9686,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -9832,15 +9756,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -9881,15 +9801,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -9916,16 +9832,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -9965,9 +9873,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -9977,12 +9883,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_deployment_usable_by_team_uses_dynamic_model_info_get(): @@ -10022,9 +9923,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -10036,21 +9935,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -10123,22 +10012,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -10326,9 +10206,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", custom_llm_provider="anthropic", model_id="model-123", @@ -10408,12 +10286,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -10438,9 +10312,7 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -10489,9 +10361,7 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -10518,9 +10388,7 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -10602,9 +10470,7 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -10635,9 +10501,7 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -10666,9 +10530,7 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -10860,12 +10722,8 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -10897,12 +10755,8 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -11108,6 +10962,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -11134,6 +10989,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -11278,7 +11134,12 @@ def _shared_ptu_deployment(model_name: str, deployment_id: str, **litellm_params """A one-PTU Azure deployment whose whole reservation is team-a's share.""" return { "model_name": model_name, - "litellm_params": {"model": "azure/gpt-4.1", "api_key": "fake", "api_base": "https://ptu.example", **litellm_params}, + "litellm_params": { + "model": "azure/gpt-4.1", + "api_key": "fake", + "api_base": "https://ptu.example", + **litellm_params, + }, "model_info": { "id": deployment_id, "base_model": "azure/gpt-4.1", @@ -11388,23 +11249,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -11449,12 +11304,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -11579,9 +11430,7 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -11626,9 +11475,7 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11703,9 +11550,7 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11772,9 +11617,7 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -11829,9 +11672,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) + router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) is None ) assert ( @@ -11853,10 +11694,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) + pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -11882,10 +11721,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) + default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None ) assert default_result is not None resolved_model, default_deployment = default_result @@ -11893,10 +11730,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) + assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" def _router_with_two_deployments(blocked_flags): @@ -11944,10 +11778,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -11981,6 +11812,7 @@ async def test_health_probe_preserves_normal_caller_policy( strict_ids: tuple[str, ...], ) -> None: import time + from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError router: Final = Router( @@ -12074,9 +11906,7 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -12085,9 +11915,7 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -12104,9 +11932,7 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -12194,9 +12020,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -12205,9 +12029,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): @@ -12249,9 +12071,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -12262,12 +12082,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -12303,7 +12118,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): missing_blocked = types.SimpleNamespace() assert litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) + litellm.Router._core_is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True ) @@ -12415,9 +12232,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -12435,22 +12250,16 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): + def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -12466,22 +12275,13 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 def test_passthrough_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) @@ -12800,9 +12600,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -12811,12 +12609,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -12831,12 +12624,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -13125,13 +12913,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -13175,12 +12963,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -13403,13 +13187,8 @@ def test_get_model_listing_info_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_reads_deployment_model_info(): @@ -13451,13 +13230,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -13516,13 +13290,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None @pytest.mark.parametrize( @@ -13551,7 +13320,10 @@ def test_get_configured_service_tiers_returns_one_value_per_deployment_in_model_ "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"service_tiers": ["ultrafast"]}, }, - {"model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}}, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}, + }, { "model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://b.example"}, @@ -13669,7 +13441,11 @@ def test_get_search_tool_support_reports_each_routable_deployment_of_the_group() @pytest.mark.parametrize("provider", ["chatgpt", "github_copilot"]) @pytest.mark.parametrize("explicit_provider", [False, True]) async def test_search_tool_catalog_lookup_needs_no_credentials( - provider: str, explicit_provider: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + provider: str, + explicit_provider: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, ) -> None: monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "chatgpt")) monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) @@ -14269,7 +14045,12 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme _AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response( 400, - json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}}, + json={ + "error": { + "code": "invalid_deployment_type", + "message": "The current deployment is not supported for batch jobs", + } + }, ) @@ -14290,7 +14071,9 @@ async def test_acreate_batch_on_a_mixed_group_reaches_chat_as_on_main_while_its_ _cool_down(router, *(d["model_info"]["id"] for d in router.model_list if d["model_info"]["mode"] == "batch")) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) with pytest.raises(openai.BadRequestError): await _create_batches(router, "gpt-mixed", attempts=1) @@ -14316,7 +14099,9 @@ async def test_acreate_batch_with_tags_that_match_only_the_chat_deployment_reach ) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) batch_route: Final = respx_mock.post(host="batch-sku.openai.azure.com").mock( return_value=httpx.Response(200, json=_AZURE_BATCH_CREATED) ) @@ -14406,13 +14191,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -14441,13 +14229,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14481,7 +14269,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -14497,13 +14287,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14512,9 +14302,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -14528,7 +14316,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -14563,9 +14353,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -14870,9 +14658,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -14981,9 +14767,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"] @@ -15170,22 +14954,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -15457,26 +15235,34 @@ class TestClaudeCodeSubagentSessionRouterBinding: @pytest.mark.asyncio async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache + import fakeredis + from redis.exceptions import ConnectionError as RedisConnectionError + + from tests._support.fake_redis import fake_redis_cache router = self._router() del router.complexity_routers["smart-router"] - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value="smart-router") - redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=3600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) - response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=self._request_kwargs(agent_id="agent-1234"), - ) + with patch.object(fakeredis.FakeAsyncRedis, "delete", side_effect=RedisConnectionError("unavailable")): + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) assert response is None - redis_cache.async_delete_cache.assert_awaited_once() + assert await raw.get(key) == b'"smart-router"' @pytest.mark.asyncio async def test_redis_read_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache + import fakeredis + + from tests._support.fake_redis import fake_redis_cache router = self._router() request_kwargs = self._request_kwargs(agent_id="agent-1234") @@ -15486,9 +15272,10 @@ class TestClaudeCodeSubagentSessionRouterBinding: cache_key, "smart-router", ) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) + server: Final = fakeredis.FakeServer() + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) + server.connected = False response = await router.async_pre_routing_hook( model="expensive-model", @@ -15497,45 +15284,46 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None assert "model_group" not in request_kwargs["metadata"] - redis_cache.async_get_cache.assert_awaited_once() @pytest.mark.asyncio async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): - from litellm.caching.caching import RedisCache + import fakeredis + from redis.exceptions import ConnectionError as RedisConnectionError + + from tests._support.fake_redis import fake_redis_cache router = self._router() - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value="smart-router") - redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=3600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) - main_response = await router.async_pre_routing_hook( - model="smart-router", - request_kwargs=self._request_kwargs(), - ) - subagent_response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=self._request_kwargs(agent_id="agent-1234"), - ) + with patch.object(fakeredis.FakeAsyncRedis, "set", side_effect=RedisConnectionError("unavailable")): + main_response = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=self._request_kwargs(), + ) + subagent_response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) assert main_response is not None assert main_response.model == "cheap-model" assert subagent_response is not None assert subagent_response.model == "cheap-model" - assert redis_cache.async_set_cache.await_count == 2 + assert await raw.get(key) == b'"smart-router"' @pytest.mark.asyncio async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): - from types import SimpleNamespace + import fakeredis - from litellm.caching.caching import RedisCache + from tests._support.fake_redis import fake_redis_cache - shared_binding = SimpleNamespace(value=None) - shared_redis = MagicMock(spec=RedisCache) - shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) - shared_redis.async_set_cache = AsyncMock( - side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) - ) + server: Final = fakeredis.FakeServer() + shared_redis: Final = fake_redis_cache(server) main_worker, subagent_worker = self._router(), self._router() main_worker.update_redis_cache(cache=shared_redis) subagent_worker.update_redis_cache(cache=shared_redis) @@ -15555,27 +15343,33 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert first.model == "cheap-model" assert second is not None assert second.model == "expensive-model" - assert shared_binding.value == "premium-router" + assert ( + await fakeredis.FakeAsyncRedis(server=server).get("claude_code_session_router:v1:key-hash-a:session-1234") + == b'"premium-router"' + ) @pytest.mark.asyncio - async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): - from litellm.caching.caching import RedisCache + @patch("time.time", return_value=1800000000.0) + async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self, _clock: object): + import fakeredis + + from tests._support.fake_redis import fake_redis_cache router = self._router() router.complexity_routers.clear() - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock() - redis_cache.async_delete_cache = AsyncMock() + server: Final = fakeredis.FakeServer() + raw: Final = fakeredis.FakeAsyncRedis(server=server) + key: Final = "claude_code_session_router:v1:key-hash-a:session-1234" + await raw.set(key, b'"smart-router"', ex=600) + redis_cache = fake_redis_cache(server) router.update_redis_cache(cache=redis_cache) for request_kwargs in (self._request_kwargs(), self._request_kwargs(agent_id="agent-1234")): response = await router.async_pre_routing_hook(model="expensive-model", request_kwargs=request_kwargs) assert response is None - redis_cache.async_get_cache.assert_not_awaited() - redis_cache.async_set_cache.assert_not_awaited() - redis_cache.async_delete_cache.assert_not_awaited() + assert await raw.get(key) == b'"smart-router"' + assert await raw.pttl(key) == 600000 @pytest.mark.asyncio async def test_session_bindings_do_not_evict_router_rate_limit_state(self): @@ -15945,18 +15739,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -16088,9 +15878,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -16364,9 +16152,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -16432,9 +16218,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -16576,9 +16360,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -16820,6 +16602,7 @@ class TestTeamPublicNameReachesPreRoutingStrategies: assert response is not None assert response.model == "gemini-flash" + @pytest.mark.asyncio async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): router = self._router( @@ -16878,7 +16661,6 @@ class TestTeamPublicNameReachesPreRoutingStrategies: with pytest.raises(litellm.BadRequestError, match="multiple teams"): two_teams._team_deployments_across_teams(self.PUBLIC_NAME) - def test_compression_policy_follows_the_same_resolution_for_every_principal(self): from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model @@ -17102,7 +16884,6 @@ class TestAutoRouterCompressionDecoupling: @pytest.mark.usefixtures("local_model_cost_map") - @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -17129,17 +16910,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -17147,17 +16925,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -17184,6 +16959,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -17293,11 +17069,15 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o def test_model_group_info_fast_mode_uses_exact_provider_catalog( local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool ) -> None: - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": operator_flag}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": operator_flag}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17309,16 +17089,21 @@ def test_model_group_info_fast_mode_uses_exact_provider_catalog( def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object ) -> None: - entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items() - if key != "supports_fast_mode"} + entry: Final = { + key: value for key, value in litellm.model_cost["claude-opus-5"].items() if key != "supports_fast_mode" + } if flag is not None: entry["supports_fast_mode"] = flag monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": True}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": True}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17326,24 +17111,30 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( assert result.supports_fast_mode is False -@pytest.mark.parametrize("other_model,expected", [ - ("anthropic/claude-opus-4-8", True), - ("anthropic/claude-opus-4-7", False), - ("anthropic/off-map-opus", False), - ("vertex_ai/claude-opus-5", False), - ("bedrock/claude-opus-5", False), -]) +@pytest.mark.parametrize( + "other_model,expected", + [ + ("anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus-4-7", False), + ("anthropic/off-map-opus", False), + ("vertex_ai/claude-opus-5", False), + ("bedrock/claude-opus-5", False), + ], +) @pytest.mark.parametrize("reverse", [True, False]) def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool ) -> None: - models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ( - "anthropic/claude-opus-5", other_model + models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ("anthropic/claude-opus-5", other_model) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "api_key": "fake-key"}, + } + for model in models + ] ) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "api_key": "fake-key"}, - } for model in models]) result: Final = router.get_model_group_info("fast-group") @@ -17583,6 +17374,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -17701,10 +17493,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -17749,9 +17538,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -17770,12 +17557,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -17787,9 +17576,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -17840,7 +17627,9 @@ async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwar await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -17885,7 +17674,9 @@ async def test_anthropic_messages_no_fallback_message_start_reaches_client_befor await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() content_released.set() @@ -17950,7 +17741,9 @@ async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecy await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -17990,8 +17783,15 @@ def _anthropic_messages_two_order_primary_model_list() -> list: id="wildcard-overridden-by-request-none", ), pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), + pytest.param( + {"fallbacks": None}, + {"model": "primary", "fallbacks": [{"model": "fallback"}]}, + True, + id="request-dict-fallback", + ), + pytest.param( + {"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback" + ), pytest.param( {"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary", "disable_fallbacks": True}, @@ -18004,7 +17804,9 @@ def _anthropic_messages_two_order_primary_model_list() -> list: True, id="content-policy-fallback", ), - pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), + pytest.param( + {"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover" + ), ], ) def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): @@ -18077,7 +17879,9 @@ async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames( await content_released.wait() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) pending = asyncio.ensure_future(wrapped.__anext__()) await asyncio.sleep(0.2) @@ -18122,7 +17926,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -18819,7 +18625,9 @@ class _AnthropicMessagesScriptedProvider: async def __call__(self, **kwargs): litellm_metadata = kwargs.get("litellm_metadata") or {} - self.calls.append((kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries"))) + self.calls.append( + (kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries")) + ) assert self._streams, "provider called more times than scripted" return self._streams.pop(0)() @@ -18837,7 +18645,9 @@ def _anthropic_messages_transport_drop(original_exception: Exception | None = No def _anthropic_messages_dropped_before_content(): - return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop()) + return _AnthropicMessagesRaisingByteStream( + [_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop() + ) def _anthropic_messages_bridge_error_chunk() -> bytes: @@ -18965,9 +18775,7 @@ async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_pa provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_logging_obj=logging_obj - ) + stream: Final = await _anthropic_messages_stream_through_router(router, provider, litellm_logging_obj=logging_obj) body: Final = tuple([chunk async for chunk in stream]) assert stash_recorder.observed == (None,) @@ -19563,20 +19371,28 @@ async def test_anthropic_messages_retries_keep_the_budget_the_first_drop_committ }, { "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-b", "api_key": "sk-test", "num_retries": sibling_num_retries}, + "litellm_params": { + "model": "anthropic/glm-b", + "api_key": "sk-test", + "num_retries": sibling_num_retries, + }, }, ], num_retries=0, fallbacks=None, ) router.set_custom_routing_strategy(_AnthropicMessagesAlternatingDeployments(router, "glm")) - provider = _AnthropicMessagesScriptedProvider(*[_anthropic_messages_dropped_before_content] * len(expected_counters)) + provider = _AnthropicMessagesScriptedProvider( + *[_anthropic_messages_dropped_before_content] * len(expected_counters) + ) stream = await _anthropic_messages_stream_through_router(router, provider) with pytest.raises(litellm.APIConnectionError): [chunk async for chunk in stream] - assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[: len(expected_counters)] + assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[ + : len(expected_counters) + ] assert [(attempted, budget) for _, attempted, budget in provider.calls] == expected_counters @@ -20911,9 +20727,7 @@ def _router_with_failing_primary(fallback_access_check) -> Router: class _DenyFallbackBudgetCheck: - async def __call__( - self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router - ) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: return False @@ -21206,9 +21020,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) + assert litellm.Router._declared_param_allowlist( + {"allowed_openai_params": ["reasoning_effort", 3]} + ) == frozenset({"reasoning_effort"}) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -21288,7 +21102,11 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], + [ + {"model_name": "x"}, + {"model_name": "x", "litellm_params": {}}, + {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, + ], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -21593,9 +21411,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -21628,9 +21444,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -22986,7 +22800,9 @@ async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit( async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): with pytest.raises(litellm.RateLimitError) as overflow: - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): + async with router._deployment_slot( + deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None + ): pass assert overflow.value.status_code == 429 assert "slot-deployment" in overflow.value.message @@ -23130,33 +22946,33 @@ def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change assert router.cached_model_group_info("grp").max_output_tokens == 200 -class _OpenBreakerRedis: - def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker - - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() - - @_redis_circuit_breaker_guard - async def async_get_cache(self, key, **kwargs): - raise AssertionError("never reached") - - @pytest.mark.asyncio async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warning(caplog): + import fakeredis + + from litellm.redis.transport.contract import RedisFailure + from tests._support.fake_redis import fake_redis_cache + router = litellm.Router( model_list=[{"model_name": "haiku", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}}] ) - router._claude_code_session_router_cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + router.update_redis_cache(cache=cache) + server.connected = False + for _ in range(6): + await cache.transport().get("probe") + result: Final = await cache.transport().get("probe") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" caplog.clear() with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - binding = await router._get_claude_code_session_router_binding("quiet-session") + binding = await router._get_claude_code_session_router_binding( + "claude_code_session_router:v1:key-hash:quiet-session" + ) assert binding is None assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] - assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) @pytest.mark.asyncio @@ -23221,18 +23037,25 @@ class TestMemberAutoRouterInference: self.cache = UserApiKeyCache() self.team = LiteLLM_TeamTable( - team_id="router-team", models=["member-router", "permitted-model"], + team_id="router-team", + models=["member-router", "permitted-model"], members_with_roles=[Member(user_id="router-member", role="user")], ) self.actor = UserAPIKeyAuth( - user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, - models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + user_id="router-member", + team_id="router-team", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], + api_key="test-key-hash", + config={"timeout": 60}, + ) + self.database = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + ) ) - self.database = SimpleNamespace(db=SimpleNamespace( - litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), - litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), - litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), - )) monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) monkeypatch.setattr(proxy_server, "prisma_client", self.database) @@ -23242,44 +23065,71 @@ class TestMemberAutoRouterInference: return { "model_name": "model_name_router-team_member-router", "litellm_params": { - "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "model": "auto_router/complexity_router", + "complexity_router_default_model": target, "complexity_router_config": { - "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), + "adaptive": False, **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), }, - "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + "tags": ["member" if member else "admin"], + "timeout": 13.0 if member else 29.0, }, "model_info": { - "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + "team_id": "router-team", + "team_public_model_name": "member-router", + "member_auto_router": member, }, } @classmethod def _router(cls, *markers: dict[str, object]) -> Router: - return Router(model_list=[ - *(markers or (cls._marker(),)), - {"model_name": "permitted-model", "litellm_params": { - "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", - }}, - {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, - ]) + return Router( + model_list=[ + *(markers or (cls._marker(),)), + { + "model_name": "permitted-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + }, + }, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ] + ) def _request( - self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + self, + *, + actor: UserAPIKeyAuth | None = None, + metadata_name: str = "metadata", + tag: str = "member", ) -> dict[str, object]: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} - if metadata_name == "litellm_metadata" else {})}, - user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + data={ + metadata_name: {"tags": [tag]}, + **( + {"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" + else {} + ), + }, + user_api_key_dict=actor or self.actor, + _metadata_variable_name=metadata_name, ) async def _route( - self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + self, + router: Router, + request: dict[str, object] | None = None, + model: str = "member-router", ) -> PreRoutingHookResponse: response: Final = await router.async_pre_routing_hook( - model=model, request_kwargs=request if request is not None else self._request(), + model=model, + request_kwargs=request if request is not None else self._request(), messages=[{"role": "user", "content": "Hello"}], ) assert response is not None @@ -23288,36 +23138,68 @@ class TestMemberAutoRouterInference: @pytest.mark.asyncio @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( - self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + self, + metadata_name: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_checks import delete_cache_team_object monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") router: Final = self._router(self._marker(classifier=True)) - classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ - "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - }) - request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", - }}} + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 200, + json={ + "id": "classifier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + request: Final = { + **self._request(metadata_name=metadata_name), + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + } + }, + } first: Final = await self._route(router, request) assert first.model == "permitted-model" and first.routing_decision is not None assert first.routing_decision["cause"] == "llm_classifier" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_teamtable.find_unique.await_count == 1 assert self.database.db.litellm_teammembership.find_unique.await_count == 1 - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={"members_with_roles": []} + ) await delete_cache_team_object( - team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + team_id=self.team.team_id, + team_alias=None, + user_api_key_cache=self.cache, + proxy_logging_obj=None, ) with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, request) - rebound: Final = {**request, "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", - }}} + rebound: Final = { + **request, + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + "x-claude-code-agent-id": "subagent", + } + }, + } with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, rebound, model="restricted-model") assert classify.call_count == 2 @@ -23327,11 +23209,23 @@ class TestMemberAutoRouterInference: async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server - request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { - "team_id": "router-team", "user_role": "proxy_admin", - }}} if state == "forged" else self._request(actor=self.actor.model_copy( - update={"user_id": ""} if state == "empty-user" else {}, - )) + request: Final = ( + { + "metadata": { + "user_api_key_team_id": "router-team", + "user_api_key_auth": { + "team_id": "router-team", + "user_role": "proxy_admin", + }, + } + } + if state == "forged" + else self._request( + actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + ) + ) + ) self.database.db.litellm_teamtable.find_unique.return_value = ( None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) ) @@ -23342,11 +23236,20 @@ class TestMemberAutoRouterInference: assert error.value.status_code == (503 if state == "unavailable" else 403) @pytest.mark.asyncio - @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) - async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: - assert (await self._route(self._router(), self._request( - actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), - ))).model == "permitted-model" + @pytest.mark.parametrize( + "user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)] + ) + async def test_service_key_and_admin_preserve_runtime_access( + self, user_id: str | None, role: LitellmUserRoles + ) -> None: + assert ( + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ), + ) + ).model == "permitted-model" @pytest.mark.asyncio @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) @@ -23357,35 +23260,56 @@ class TestMemberAutoRouterInference: from litellm.proxy._types import LiteLLM_ProjectTableCachedObj from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if ceiling == "team" else self.team.models, - "organization_id": "router-org" if ceiling == "organization" else None, - }) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + } + ) if ceiling == "member": await self.cache.async_set_cache( key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), - value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + value=LiteLLM_TeamMembership( + user_id="router-member", + team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"]), + ), model_type=LiteLLM_TeamMembership, ) elif ceiling == "organization": await self.cache.async_set_cache( - key="org_id:router-org", value=LiteLLM_OrganizationTable( - organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + key="org_id:router-org", + value=LiteLLM_OrganizationTable( + organization_id="router-org", + budget_id="org-budget", + created_by="admin", + updated_by="admin", models=["restricted-model"], - ), model_type=LiteLLM_OrganizationTable, + ), + model_type=LiteLLM_OrganizationTable, ) elif ceiling == "project": await self.cache.async_set_cache( - key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( - project_id="router-project", team_id="router-team", models=["restricted-model"], - ), model_type=LiteLLM_ProjectTableCachedObj, + key="project_id:router-project", + value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", + team_id="router-team", + models=["restricted-model"], + ), + model_type=LiteLLM_ProjectTableCachedObj, ) with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if ceiling == "key" else self.actor.models, - "project_id": "router-project" if ceiling == "project" else None, - }))) + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + } + ) + ), + ) @pytest.mark.asyncio @pytest.mark.parametrize("group_owner", ("team", "key")) @@ -23393,22 +23317,32 @@ class TestMemberAutoRouterInference: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast group: Final = LiteLLM_AccessGroupTable( - access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + access_group_id="router-group", + access_group_name="Router targets", + access_model_names=["permitted-model"], ) self.database.db.litellm_accessgrouptable.find_unique.return_value = group - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if group_owner == "team" else self.team.models, - "access_group_ids": ["router-group"] if group_owner == "team" else [], - }) - request: Final = self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if group_owner == "key" else self.actor.models, - "access_group_ids": ["router-group"] if group_owner == "key" else [], - })) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + } + ) + request: Final = self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + } + ) + ) router: Final = self._router() assert (await self._route(router, request)).model == "permitted-model" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 - self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy( + update={"access_model_names": []} + ) await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) with pytest.raises(ProxyException, match="is not available for this API key"): await self._route(router, request) @@ -23419,13 +23353,16 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False), self._marker()) request: Final = self._request() selected: Final = router._selected_strategy_marker_deployment( - model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, + model="model_name_router-team_member-router", + strategy_tags=("member",), + request_kwargs=request, ) assert selected is not None and selected["model_info"]["member_auto_router"] is True assert (await self._route(router, request)).model == "permitted-model" assert request["timeout"] == 13.0 await self.cache.async_set_cache( - key="team_id:router-team", model_type=LiteLLM_TeamTable, + key="team_id:router-team", + model_type=LiteLLM_TeamTable, value=self.team.model_copy(update={"models": ["member-router"]}), ) with pytest.raises(ProxyException, match="is not available for this API key"): @@ -23441,7 +23378,9 @@ class TestMemberAutoRouterInference: router: Final = self._router(self._marker(member=False)) monkeypatch.setitem(sys.modules, "fastapi", None) monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) - assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" + assert ( + await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}}) + ).model == "restricted-model" def _access_window_offsets(start_hours: float, end_hours: float, team_ids: list) -> dict: @@ -23486,10 +23425,12 @@ def test_access_windows_hide_reserved_deployment_from_other_teams(): def test_access_windows_raise_when_only_reserved_deployments_remain(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], - )[:1]) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], + )[:1] + ) for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-b"}}, {}): with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment(model="gpt-4o-ptu", request_kwargs=request_kwargs) @@ -23552,9 +23493,11 @@ def test_access_windows_inactive_window_leaves_deployments_available(): def test_access_windows_apply_when_calling_by_model_id(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - )) + router = Router( + model_list=_reserved_model_list( + windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], + ) + ) with pytest.raises(litellm.BadRequestError, match="reserved for another team"): router._common_checks_available_deployment( model="reserved-deployment", @@ -23800,7 +23743,9 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref @pytest.mark.asyncio -async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> None: +async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> ( + None +): """The context-window chain is keyed the same way the ordinary chain is, so a key spelled like the wildcard deployment ("anthropic/claude-sonnet-4-6") must catch the bare group's context-window error too.""" router = litellm.Router( @@ -23982,9 +23927,7 @@ def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the [ GuardrailRaisedException(guardrail_name="chunk-scanner", message="blocked"), HTTPException(status_code=403, detail={"error": "blocked", "guardrail_name": "chunk-scanner"}), - ModifyResponseException( - message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner" - ), + ModifyResponseException(message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner"), ], ) async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: Exception) -> None: @@ -24185,6 +24128,7 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(): assert seen == ["router_usage"] assert current_service_target() is None + class _SpanRecordingInMemoryCache(InMemoryCache): """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in.""" @@ -24379,10 +24323,13 @@ async def test_non_chat_surfaces_mark_their_deployment_pick(monkeypatch: pytest. router.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) assert events == [_pick("embed", "initial", 1), _pick("gpt-4o", "initial", 1)] + + class TestAutoRouterTraceProvenance: @pytest.fixture def _tracing(self, monkeypatch: pytest.MonkeyPatch) -> "tuple[OpenTelemetryV2, InMemorySpanExporter]": from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from litellm.integrations.otel import OpenTelemetryV2Config from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.plumbing import providers @@ -24390,7 +24337,9 @@ class TestAutoRouterTraceProvenance: config: Final = OpenTelemetryV2Config(exporter="in_memory") exporter: Final = InMemorySpanExporter() - logger: Final = OpenTelemetryV2(config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter)) + logger: Final = OpenTelemetryV2( + config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter) + ) monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) return logger, exporter @@ -24400,7 +24349,8 @@ class TestAutoRouterTraceProvenance: "model_name": "traced-router", "model_info": {"id": f"definition-{tag}", "updated_at": "2026-10-01T12:00:00+00:00"}, "litellm_params": { - "model": "auto_router/complexity_router", "tags": [tag], + "model": "auto_router/complexity_router", + "tags": [tag], "complexity_router_config": { "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), }, @@ -24409,17 +24359,25 @@ class TestAutoRouterTraceProvenance: def test_definition_fingerprint_is_canonical_and_tracks_config_not_credentials(self): original: Final = Deployment.model_validate(self._marker("blue")) - reordered: Final = Deployment.model_validate({ - **self._marker("blue"), - "litellm_params": { - "api_key": "credential-only-change", "tags": ["blue"], "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer")}, - }, - }) + reordered: Final = Deployment.model_validate( + { + **self._marker("blue"), + "litellm_params": { + "api_key": "credential-only-change", + "tags": ["blue"], + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer") + }, + }, + } + ) fingerprint: Final = Router._routing_definition_fingerprint(original) assert fingerprint is not None assert fingerprint == Router._routing_definition_fingerprint(reordered) - assert fingerprint != Router._routing_definition_fingerprint(Deployment.model_validate(self._marker("blue", "other"))) + assert fingerprint != Router._routing_definition_fingerprint( + Deployment.model_validate(self._marker("blue", "other")) + ) @pytest.mark.asyncio async def test_direct_classifier_and_plugin_prelude_never_inherit_parent_decisions( @@ -24437,7 +24395,9 @@ class TestAutoRouterTraceProvenance: async def classify(self, context: RoutingContext) -> str: await litellm.acompletion( - model="openai/gpt-5.6-luna", api_key="test", mock_response="SIMPLE", + model="openai/gpt-5.6-luna", + api_key="test", + mock_response="SIMPLE", messages=[{"role": "user", "content": "choose a tier"}], metadata={**context.metadata, "internal_call_origin": "autorouter_classifier"}, ) @@ -24451,16 +24411,25 @@ class TestAutoRouterTraceProvenance: monkeypatch.setattr(litellm, "_async_success_callback", [logger, CallbackComplete()]) classifier: Final = ObservedClassifier() - router: Final = Router(model_list=[ - {**self._marker("blue"), "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "classifier_type": "custom", "classifier_plugin": classifier, - "tiers": {"SIMPLE": "answer"}, "session_affinity": False, "deployment_affinity": False, + router: Final = Router( + model_list=[ + { + **self._marker("blue"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "custom", + "classifier_plugin": classifier, + "tiers": {"SIMPLE": "answer"}, + "session_affinity": False, + "deployment_affinity": False, + }, + }, }, - }}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], plugins=[classifier]) + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + plugins=[classifier], + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale, "user_api_key_team_id": "team"}} with logger.tracer.start_as_current_span("request") as root: @@ -24470,7 +24439,8 @@ class TestAutoRouterTraceProvenance: ) await completed.wait() classifier_span: Final = next( - span for span in exporter.get_finished_spans() + span + for span in exporter.get_finished_spans() if (span.attributes or {}).get("litellm.request.purpose") == "autorouter_classifier" ) assert not any(key.startswith("litellm.routing.") for key in classifier_span.attributes) @@ -24484,36 +24454,66 @@ class TestAutoRouterTraceProvenance: @pytest.mark.asyncio @pytest.mark.parametrize("build_fails", [False, True]) async def test_semantic_internal_failure_keeps_phase_identity_without_parent_retry_attribution( - self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, build_fails: bool, + self, + _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + build_fails: bool, ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) query_statuses: Final = iter((503, 200)) def embedding_response(request: httpx.Request) -> httpx.Response: - status: Final = 400 if build_fails else ( - next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200 + status: Final = ( + 400 + if build_fails + else (next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200) + ) + return httpx.Response( + status, + headers={"x-should-retry": "false"}, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, ) - return httpx.Response(status, headers={"x-should-retry": "false"}, json={ - "object": "list", "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock(side_effect=embedding_response) - router: Final = Router(model_list=[ - {"model_name": "semantic", "model_info": {"id": "semantic-definition"}, "litellm_params": { - "model": "auto_router/semantic-router", "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - }}, - *({"model_name": "embed", "model_info": {"id": f"embed-{index}"}, "litellm_params": { - "model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", - "api_key": "test", "max_retries": 0, - }} for index in range(2)), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], num_retries=1, retry_after=0) + upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock( + side_effect=embedding_response + ) + router: Final = Router( + model_list=[ + { + "model_name": "semantic", + "model_info": {"id": "semantic-definition"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + }, + }, + *( + { + "model_name": "embed", + "model_info": {"id": f"embed-{index}"}, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + "max_retries": 0, + }, + } + for index in range(2) + ), + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + num_retries=1, + retry_after=0, + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale}} with logger.tracer.start_as_current_span("request") as root: @@ -24545,20 +24545,22 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]" ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), - { - "model_name": "traced-router", - "model_info": {"id": "definition-semantic"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "tags": ["blue"], - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", + router: Final = Router( + model_list=[ + self._marker("blue"), + { + "model_name": "traced-router", + "model_info": {"id": "definition-semantic"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "tags": ["blue"], + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, }, - }, - ]) + ] + ) with logger.start_phase_span("route traced-router"): response: Final = await router.async_pre_routing_hook( model="traced-router", request_kwargs={"metadata": {"tags": ["blue"]}} @@ -24578,26 +24580,42 @@ class TestAutoRouterTraceProvenance: ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - respx_mock.post("https://embedding.example/v1/embeddings").respond(200, json={ - "object": "list", - "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - router: Final = Router(model_list=[ - { - "model_name": "internal-semantic-router", - "model_info": {"id": "semantic-definition", "team_id": "trace-team", "team_public_model_name": "public-router"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - }, + respx_mock.post("https://embedding.example/v1/embeddings").respond( + 200, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, }, - {"model_name": "embed", "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", "api_key": "test"}}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ]) + ) + router: Final = Router( + model_list=[ + { + "model_name": "internal-semantic-router", + "model_info": { + "id": "semantic-definition", + "team_id": "trace-team", + "team_public_model_name": "public-router", + }, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + }, + }, + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ] + ) request: Final = {"metadata": {"user_api_key_team_id": "trace-team"}} with logger.tracer.start_as_current_span("request"): await router.async_get_available_deployment( @@ -24609,8 +24627,13 @@ class TestAutoRouterTraceProvenance: assert decision["router_config_id"] == router.model_list[0]["model_info"]["id"] assert decision["cause"] == "semantic_match" spans: Final = {span.name: span for span in exporter.get_finished_spans()} - assert spans["route public-router"].attributes["litellm.routing.router_model_name"] == decision["router_model_name"] - assert spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + assert ( + spans["route public-router"].attributes["litellm.routing.router_model_name"] + == decision["router_model_name"] + ) + assert ( + spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + ) retry: Final = next(event for event in spans["request"].events if event.name == "litellm.routing.retry") assert retry.attributes["litellm.routing.router_model_name"] == decision["router_model_name"] @@ -24620,18 +24643,37 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", passthrough: bool ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), self._marker("green"), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test", "use_in_pass_through": True, "tags": ["blue", "green"]}}, - ], enable_tag_filtering=True) + router: Final = Router( + model_list=[ + self._marker("blue"), + self._marker("green"), + { + "model_name": "answer", + "litellm_params": { + "model": "openai/gpt-5.6-luna", + "api_key": "test", + "use_in_pass_through": True, + "tags": ["blue", "green"], + }, + }, + ], + enable_tag_filtering=True, + ) request: Final = {"metadata": {"tags": ["green"]}, "litellm_metadata": {"tags": ["green"]}} - pick: Final = router.async_get_available_deployment_for_pass_through if passthrough else router.async_get_available_deployment + pick: Final = ( + router.async_get_available_deployment_for_pass_through + if passthrough + else router.async_get_available_deployment + ) with logger.tracer.start_as_current_span("request"): await pick(model="traced-router", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) decision: Final = request["litellm_metadata"]["routing_decision"] assert decision["router_config_id"] == "definition-green" assert decision["router_config_updated_at"] == "2026-10-01T12:00:00+00:00" - assert decision["router_config_fingerprint"] == router.complexity_routers["traced-router"][1].definition_fingerprint + assert ( + decision["router_config_fingerprint"] + == router.complexity_routers["traced-router"][1].definition_fingerprint + ) assert "routing_decision" not in request["metadata"] router.log_retry({**request, "model": "answer"}, ValueError("private-error-text")) await pick(model="answer", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) @@ -24652,6 +24694,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -24704,6 +24747,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -24724,6 +24768,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -24746,6 +24791,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @patch("azure.identity.get_bearer_token_provider") @patch("azure.identity.ClientSecretCredential") @@ -24827,6 +24873,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab # finally verify if the mocked token was used by Azure SDK mocked_func_generating_token.assert_called() + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_audio_speech_router(): @@ -24870,6 +24917,7 @@ async def test_audio_speech_router(): client_passed_in_request = mock_aspeech.call_args.kwargs["client"] assert client_passed_in_request == expected_openai_client + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_update_kwargs_before_fallbacks_unit_test(): router = Router( @@ -24895,6 +24943,7 @@ def test_update_kwargs_before_fallbacks_unit_test(): assert kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "call_type", @@ -24954,6 +25003,7 @@ async def test_update_kwargs_before_fallbacks(call_type): print(mock_client.call_args.kwargs) assert mock_client.call_args.kwargs["litellm_trace_id"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_model_info_wildcard_routes(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -24973,6 +25023,7 @@ def test_router_get_model_info_wildcard_routes(): assert model_info["tpm"] is not None assert model_info["rpm"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) @@ -25003,6 +25054,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): assert tpm is not None, "tpm is None" assert rpm is not None, "rpm is None" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_call_router_callbacks_on_success(): @@ -25036,6 +25088,7 @@ async def test_call_router_callbacks_on_success(): assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") assert increment["increment_value"] == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio @@ -25064,6 +25117,7 @@ async def test_call_router_callbacks_on_failure(): assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") + @pytest.mark.asyncio async def test_rate_limit_error_callback(): """ @@ -25109,13 +25163,8 @@ async def test_rate_limit_error_callback(): except Exception: pass - with patch.object( - customHandler, "log_model_group_rate_limit_error", new=AsyncMock() - ) as mock_client: - - print( - f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}" - ) + with patch.object(customHandler, "log_model_group_rate_limit_error", new=AsyncMock()) as mock_client: + print(f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}") try: _ = await router.acompletion( @@ -25138,7 +25187,6 @@ async def test_rate_limit_error_callback(): async def test_router_model_group_headers(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - from litellm.types.utils import OPENAI_RESPONSE_HEADERS router = Router( model_list=[ @@ -25163,12 +25211,12 @@ async def test_router_model_group_headers(): assert "x-ratelimit-remaining-requests" in resp._hidden_params["additional_headers"] assert "x-ratelimit-remaining-tokens" in resp._hidden_params["additional_headers"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_remaining_model_group_usage(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - from litellm.types.utils import OPENAI_RESPONSE_HEADERS router = Router( model_list=[ @@ -25194,6 +25242,7 @@ async def test_get_remaining_model_group_usage(): assert "x-ratelimit-remaining-requests" in remaining_usage assert "x-ratelimit-remaining-tokens" in remaining_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "potential_access_group, expected_result", @@ -25212,6 +25261,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) access_groups = router.is_model_access_group_for_wildcard_route(model_access_group=potential_access_group) assert access_groups == expected_result + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) @@ -25222,6 +25272,7 @@ def test_router_redis_cache(): assert router.cache.redis_cache == redis_cache + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_handle_clientside_credential(): """A caller-supplied credential must stay scoped to the current call: it must @@ -25250,6 +25301,7 @@ def test_router_handle_clientside_credential(): assert len(router.get_model_list()) == 1 assert router.get_deployment(model_id=new_deployment.model_info.id) is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_clientside_credential_not_reused_by_other_callers(respx_mock, monkeypatch: pytest.MonkeyPatch): """End-to-end regression test for LIT-7811. @@ -25307,6 +25359,7 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router( @@ -25323,6 +25376,7 @@ def test_router_get_async_openai_model_client(): model_client = router._get_async_openai_model_client(deployment=MagicMock(), kwargs={}) assert model_client is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials(): router = Router( @@ -25338,6 +25392,7 @@ def test_router_get_deployment_credentials(): assert credentials is not None assert credentials["api_key"] == "123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider(): """ @@ -25382,6 +25437,7 @@ def test_router_get_deployment_credentials_with_provider(): credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider_wildcard(): """ @@ -25429,6 +25485,7 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_model_info(): router = Router( @@ -25443,12 +25500,14 @@ def test_router_get_deployment_model_info(): model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None + @pytest.fixture() def _vcr_outcome_gate_router_unit(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown_router_unit(): """ @@ -25470,6 +25529,7 @@ def setup_and_teardown_router_unit(): loop.close() asyncio.set_event_loop(None) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_acompletion_deployment_not_mutated(): @@ -25519,6 +25579,7 @@ async def test_acompletion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_completion_deployment_not_mutated(): """ @@ -25567,6 +25628,7 @@ def test_completion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_default_deployment_isolation(): """ @@ -25639,10 +25701,12 @@ def test_default_deployment_isolation(): assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore + class NoItemsAliasDict(dict): def items(self): raise AssertionError("Unexpected full alias iteration via items()") + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): router = Router( @@ -25659,6 +25723,7 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini") assert model_alias_list == [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): router = Router( @@ -25887,9 +25952,11 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + if __name__ == "__main__": pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_is_prompt_management_model_optimization(): """ @@ -25944,6 +26011,7 @@ def test_is_prompt_management_model_optimization(): finally: litellm._known_custom_logger_compatible_callbacks = original_callbacks + @pytest.fixture def router(): """Create a router with a mock deployment""" @@ -25959,6 +26027,7 @@ def router(): ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch(router): @@ -25982,6 +26051,7 @@ async def test_router_acancel_batch(router): assert response.id == "batch_123" assert response.status == "cancelled" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_resolves_credential_name(): @@ -26022,6 +26092,7 @@ async def test_router_acancel_batch_resolves_credential_name(): finally: litellm.credential_list = [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_removes_unresolved_credential_name(): @@ -26054,6 +26125,7 @@ async def test_router_acancel_batch_removes_unresolved_credential_name(): call_kwargs = mock_cancel.call_args.kwargs assert "litellm_credential_name" not in call_kwargs + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingHeaders: """Test that embedding methods properly propagate headers from router configuration.""" @@ -26368,6 +26440,7 @@ class TestRouterEmbeddingHeaders: assert "litellm_trace_id" in completion_kwargs assert "litellm_trace_id" in embedding_kwargs + if __name__ == "__main__": # Run a simple test test = TestRouterEmbeddingHeaders() @@ -26384,6 +26457,7 @@ OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: return respx_mock.post(url).mock( return_value=httpx.Response( @@ -26397,11 +26471,13 @@ def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route ) ) + def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: request = route.calls[index].request body = json.loads(request.read()) return request.headers["authorization"], body["model"], body["input"] + def _alias_router() -> Router: return Router( model_list=[ @@ -26415,6 +26491,7 @@ def _alias_router() -> Router: ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" @@ -26901,10 +26978,12 @@ class TestRouterEmbeddingIntegration: assert "headers" in call_kwargs assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + if __name__ == "__main__": # Run tests pytest.main([__file__, "-v"]) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterIndexManagement: """Test cases for router index management functions""" @@ -27665,7 +27744,6 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27707,7 +27785,7 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() assert response.choices[0].message.content == "Hello world 1!" except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") @@ -27719,7 +27797,6 @@ def test_router_context_window_check_pre_call_check(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27761,7 +27838,7 @@ def test_router_context_window_check_pre_call_check(): assert response.choices[0].message.content == "Hello world 2!" except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") @@ -27773,7 +27850,6 @@ def test_router_context_window_check_pre_call_check_out_group(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -27799,7 +27875,13 @@ def test_router_context_window_check_pre_call_check_out_group(): }, ] - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0, context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}]) # type: ignore + router = Router( + model_list=model_list, + set_verbose=True, + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}], + ) # type: ignore response = router.completion( model="gpt-3.5-turbo-small", @@ -27812,7 +27894,7 @@ def test_router_context_window_check_pre_call_check_out_group(): print(f"response: {response}") assert response.choices[0].message.content == "Alexander was a great conqueror." except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") def test_filter_invalid_params_pre_call_check(): @@ -27850,7 +27932,7 @@ def test_filter_invalid_params_pre_call_check(): ) assert len(filtered_deployments) == 1 except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.mark.parametrize("allowed_model_region", ["eu", None, "us"]) @@ -27895,13 +27977,9 @@ def test_router_region_pre_call_check(allowed_model_region): if allowed_model_region is None: assert len(_healthy_deployments) == 2 else: - assert len(_healthy_deployments) == 1, "{} models selected as healthy".format( - len(_healthy_deployments) - ) - assert ( - _healthy_deployments[0]["model_info"]["id"] == "1" - ), "Incorrect model id picked. Got id={}, expected id=1".format( - _healthy_deployments[0]["model_info"]["id"] + assert len(_healthy_deployments) == 1, f"{len(_healthy_deployments)} models selected as healthy" + assert _healthy_deployments[0]["model_info"]["id"] == "1", ( + "Incorrect model id picked. Got id={}, expected id=1".format(_healthy_deployments[0]["model_info"]["id"]) ) @@ -28072,14 +28150,12 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): if llm_provider == "azure" and base_model is None: assert len(updated_list) == 1 else: - pytest.fail("Expected to raise an error. Got={}".format(updated_list)) - except Exception as e: - if ( - llm_provider == "azure" and base_model is not None - ) or llm_provider == "openai": + pytest.fail(f"Expected to raise an error. Got={updated_list}") + except Exception: + if (llm_provider == "azure" and base_model is not None) or llm_provider == "openai": pass except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + pytest.fail(f"Got unexpected exception on router! - {e!s}") @pytest.fixture @@ -28153,9 +28229,7 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): ): new_retry_after_mock_client = MagicMock(return_value=-1) - litellm.utils._get_retry_after_from_exception_header = ( - new_retry_after_mock_client - ) + litellm.utils._get_retry_after_from_exception_header = new_retry_after_mock_client try: router.embedding( @@ -28241,9 +28315,7 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): ) if sync_mode: - cooldown_deployments = get_cooldown_deployments( - litellm_router_instance=router, parent_otel_span=None - ) + cooldown_deployments = get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) else: cooldown_deployments = await async_get_cooldown_deployments( litellm_router_instance=router, parent_otel_span=None @@ -28642,8 +28714,8 @@ class CompletionCustomHandler( # Class variables or attributes def __init__(self): self.errors = [] - self.states: Optional[ - List[ + self.states: ( + list[ Literal[ "sync_pre_api_call", "async_pre_api_call", @@ -28656,7 +28728,8 @@ class CompletionCustomHandler( "async_failure", ] ] - ] = [] + | None + ) = [] def log_pre_api_call(self, model, messages, kwargs): try: @@ -28680,14 +28753,10 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) - except Exception as e: + except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -28711,9 +28780,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs["original_response"]) or inspect.isasyncgen(kwargs["original_response"]) ) @@ -28725,12 +28792,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -28747,23 +28810,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -28784,22 +28842,17 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponse) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) ) + assert isinstance(kwargs["api_key"], (str, type(None))) + assert isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str) assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) @@ -28818,23 +28871,18 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -28849,8 +28897,7 @@ class CompletionCustomHandler( No-op. Not implemented yet. """ - pass - except Exception as e: + except Exception: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -28867,29 +28914,19 @@ class CompletionCustomHandler( ## END TIME assert isinstance(end_time, datetime) ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) + assert isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) ## KWARGS assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) + base_model = litellm.utils.get_base_model_from_metadata(model_call_details=kwargs) - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): + if kwargs["model"] == "chatgpt-v-3" and base_model is not None and kwargs["stream"] != True: # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float) response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) + print(f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}") prompt_tokens = response_obj.usage.prompt_tokens completion_tokens = response_obj.usage.completion_tokens # ensure the pricing is based on the base_model here @@ -28900,9 +28937,9 @@ class CompletionCustomHandler( ) expected_price = prompt_price + completion_price print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + assert response_cost == expected_price, ( + f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + ) assert isinstance(kwargs["messages"], list) assert isinstance(kwargs["optional_params"], dict) @@ -28913,9 +28950,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -28928,12 +28963,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -28960,9 +28991,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -29105,7 +29134,6 @@ class SpecialEndpointRoutingStrategy(CustomRoutingStrategyBase): if isinstance(model, dict): if model["litellm_params"]["model"] == "openai/very-special-endpoint": return model - pass def get_available_deployment( self, @@ -29238,9 +29266,7 @@ def test_multiple_deployments_route_to_a_configured_model( }, ) - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - side_effect=_deployment_response - ) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_deployment_response) router: Final = Router( model_list=[ { @@ -29389,9 +29415,7 @@ def test_router_retries_transient_server_error( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(500, json={"error": {"message": "temporary failure"}}), _openai_chat_response("recovered"), @@ -29424,9 +29448,7 @@ def test_router_authentication_error_is_not_swallowed( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 401, @@ -29524,9 +29546,9 @@ def test_router_reads_api_key_from_model_list( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(side_effect=[_openai_chat_response(), _openai_chat_stream_response()]) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[_openai_chat_response(), _openai_chat_stream_response()] + ) router: Final = Router( model_list=[ { @@ -29549,9 +29571,7 @@ def test_router_reads_api_key_from_model_list( messages=[{"role": "user", "content": "stream key from model list"}], stream=True, ) - streamed_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in stream - ) + streamed_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream) assert response.choices[0].message.content == "scripted reply" assert streamed_content == "streamed reply" @@ -29567,9 +29587,9 @@ def test_router_calls_requested_specific_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(return_value=_openai_chat_response()) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_openai_chat_response() + ) router: Final = Router( model_list=[ { @@ -29611,9 +29631,7 @@ async def test_router_context_window_error_uses_configured_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 400, @@ -29656,9 +29674,7 @@ async def test_router_context_window_error_uses_configured_fallback( assert response.choices[0].message.content == "fallback response" assert route.call_count == 2 - assert tuple( - json.loads(call.request.content)["model"] for call in route.calls - ) == ("gpt-4o-mini", "gpt-4.1-mini") + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("gpt-4o-mini", "gpt-4.1-mini") def test_router_function_calling_sends_tools_and_returns_tool_call( @@ -29666,9 +29682,7 @@ def test_router_function_calling_sends_tools_and_returns_tool_call( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -29747,9 +29761,7 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ _openai_chat_response("async reply"), _openai_chat_response("sync reply"), @@ -29784,17 +29796,13 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( stream=True, ) async_stream_chunks: Final = tuple([chunk async for chunk in async_stream]) - async_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in async_stream_chunks - ) + async_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in async_stream_chunks) sync_stream: Final = router.completion( model="openai-model", messages=[{"role": "user", "content": "sync stream"}], stream=True, ) - sync_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in sync_stream - ) + sync_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in sync_stream) assert async_response.choices[0].message.content == "async reply" assert sync_response.choices[0].message.content == "sync reply" @@ -29838,9 +29846,7 @@ async def test_router_moderation_returns_provider_result( num_retries=0, ) - response: Final = await router.amoderation( - model="moderation-model", input="safe scripted text" - ) + response: Final = await router.amoderation(model="moderation-model", input="safe scripted text") assert response.results[0].flagged is False assert json.loads(route.calls[0].request.content)["input"] == "safe scripted text" @@ -30009,9 +30015,7 @@ async def test_batch_completion_multiple_models_returns_each_model_response( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30034,9 +30038,7 @@ async def test_batch_completion_multiple_models_and_messages_preserves_each_pair respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30117,9 +30119,7 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None ) chunks: Final = tuple([chunk async for chunk in response]) text: Final = "".join( - chunk.choices[0].delta.content or "" - for chunk in chunks - if chunk.choices and chunk.choices[0].delta is not None + chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices and chunk.choices[0].delta is not None ) assert text == "stream answer" @@ -30275,7 +30275,9 @@ def _first_chunk_cases(): from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) - provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m") + provider_error = litellm.InternalServerError( + message="primary died at stream start", llm_provider="openai", model="m" + ) wrapped = MidStreamFallbackError( message=str(provider_error), model="m", @@ -30561,7 +30563,9 @@ def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_ assert [d["model_info"]["id"] for d in filtered] == ["only"] -def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -30584,7 +30588,9 @@ def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supp assert filtered == [] -def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) router: Final = Router( model_list=[ @@ -30653,7 +30659,9 @@ def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] -def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -30898,7 +30906,9 @@ def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litell Router()._update_kwargs_with_deployment( deployment={ "model_name": "gpt-5-mini", - "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "litellm_params": LiteLLM_Params( + model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options + ), "model_info": {"id": "pydantic-params"}, }, kwargs=kwargs, @@ -31296,7 +31306,10 @@ _IMAGE_RESPONSES_INPUT = [ ], } ] -_VISION_LEG: Final[dict[str, object]] = {"litellm_params": {"model": "openai/gpt-5.6-terra"}, "model_info": {"id": "vision"}} +_VISION_LEG: Final[dict[str, object]] = { + "litellm_params": {"model": "openai/gpt-5.6-terra"}, + "model_info": {"id": "vision"}, +} _DECLARED_TEXT_LEG: Final[dict[str, object]] = { "litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"id": "declared-text", "supports_vision": False}, @@ -31359,7 +31372,10 @@ def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooli def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): - assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG + assert ( + litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) + is _DECLARED_TEXT_LEG + ) def test_image_routing_drop_declared_text_only_legs_for_image_requests(): @@ -31412,16 +31428,29 @@ def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): ("deployment", "expected"), [ (_DECLARED_TEXT_LEG, True), - ({"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, False), + ( + {"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, + False, + ), ({"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {}}, True), ({"litellm_params": {"model": "azure/my-deploy"}, "model_info": {"base_model": "azure/gpt-audio"}}, True), ({"litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct"}, "model_info": {}}, False), ({"litellm_params": {"model": "no-such-model-anywhere"}, "model_info": {}}, False), ({"litellm_params": {}, "model_info": {}}, False), ], - ids=["override-false", "override-true", "cost-map-false", "base-model-false", "undeclared", "unresolvable", "no-params"], + ids=[ + "override-false", + "override-true", + "cost-map-false", + "base-model-false", + "undeclared", + "unresolvable", + "no-params", + ], ) -def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map(deployment: dict[str, object], expected: bool): +def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map( + deployment: dict[str, object], expected: bool +): assert litellm.Router._deployment_declares_no_vision(deployment, "mixed") is expected @@ -31453,19 +31482,31 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( headers={"retry-after": "168307"}, json={"type": "error", "error": {"type": "rate_limit_error", "message": "usage limit"}}, ) - return httpx.Response(200, json={ - "id": "msg-healthy", "type": "message", "role": "assistant", "model": "claude-sonnet-5-5", - "content": [{"type": "text", "text": "healthy account"}], "stop_reason": "end_turn", - "usage": {"input_tokens": 2, "output_tokens": 2}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-healthy", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "healthy account"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 2}, + }, + ) route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) router: Final = Router( - model_list=[{ - "model_name": "claude", "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", "api_key": "fake-gateway-key", - }, "model_info": {"id": "shared-anthropic"}, - }], + model_list=[ + { + "model_name": "claude", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "fake-gateway-key", + }, + "model_info": {"id": "shared-anthropic"}, + } + ], num_retries=0, allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), ) @@ -31473,22 +31514,36 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( async def call(account: str | None) -> object: credentials: Final = ( - {"provider_specific_header": { - "custom_llm_provider": "anthropic", "extra_headers": { - "Authorization": f"Bearer sk-ant-oat-fake-{account}", - }, - }} if credential_kind == "oauth" else {"api_key": f"fake-{account}"} - ) if account is not None else {} + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": { + "Authorization": f"Bearer sk-ant-oat-fake-{account}", + }, + } + } + if credential_kind == "oauth" + else {"api_key": f"fake-{account}"} + ) + if account is not None + else {} + ) if surface == "messages": max_tokens: Final = litellm.get_max_tokens("anthropic/claude-sonnet-5-5") assert max_tokens is not None return await router.aanthropic_messages( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, max_tokens=max_tokens, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + max_tokens=max_tokens, + **credentials, ) return await router.acompletion( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + **credentials, ) try: @@ -31726,7 +31781,9 @@ async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_ served: Final = component_embeddings[str(request["InferenceComponentName"])] return { "Body": io.BytesIO( - json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode() + json.dumps( + {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]} + ).encode() ) } @@ -31759,7 +31816,9 @@ async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_cu async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone( respx_mock: respx.MockRouter, _router_httpx_transport: None ) -> None: - route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY)) + route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond( + json=dict(_CHAT_COMPLETION_REPLY) + ) router: Final = Router( model_list=[ { @@ -31774,7 +31833,9 @@ async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clien ] ) - await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value") + await router.acompletion( + model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value" + ) assert json.loads(route.calls.last.request.content)["model_id"] == "client-value"