fix(sdk): stop requiring redis to import the Messages streaming path (#51685)

* fix(sdk): stop requiring redis to import the Messages streaming path

redis_update_buffer.py and dynamic_rate_limiter_v3.py imported RedisError at module
level. Both sit on the import chain that the SDK /v1/messages stream loads, so an SDK
install without redis lost every streaming success callback. Import RedisError only
where a Redis call can actually raise it.

* style(tests): space the optional redis import fixture per ruff format

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>

litellm/proxy/db/db_transaction_queue/redis_update_buffer.py

@@ -13,7 +13,6 @@ from itertools import chain
 from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
 
 from pydantic import TypeAdapter
-from redis.exceptions import RedisError
 
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
@@ -601,6 +600,7 @@ class RedisUpdateBuffer:
         )
         if len(rpush_list) == 0:
             return
+        from redis.exceptions import RedisError
 
         try:
             await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
@@ -620,6 +620,8 @@ class RedisUpdateBuffer:
         """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable."""
         if self.redis_cache is None:
             return
+        from redis.exceptions import RedisError
+
         items: Final = tool_discovery_queue.flush()
         if not items:
             return

litellm/proxy/hooks/dynamic_rate_limiter_v3.py

@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Final, Literal
 
 from fastapi import HTTPException
 from pydantic import BaseModel, ConfigDict, ValidationError
-from redis.exceptions import RedisError
 
 import litellm
 from litellm import Router
@@ -526,6 +525,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
         member: Final = (
             pool_member(priority_class) if user_api_key_dict.team_id is None else team_member(user_api_key_dict.team_id)
         )
+        from redis.exceptions import RedisError
+
         try:
             members: Final = await self.active_teams.touch(model, member, priority_class, self.v3_limiter.window_size)
         except (RedisError, RedisCircuitBreakerOpenError, ValidationError) as e:

tests/unit/proxy/conftest.py

@@ -4,13 +4,16 @@ import _hashlib
 import asyncio
 import copy
 import hashlib
+import importlib
 import inspect
 import logging
 import os
+import sys
 import tempfile
 import warnings
 from collections.abc import Awaitable, Callable, Iterator
 from contextlib import contextmanager
+from types import ModuleType
 from typing import Dict, Final, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -570,3 +573,18 @@ def free_and_paid_router() -> "litellm.Router":
             _per_token_priced("paid-model", "openai/paid-backend", 1e-6, 2e-6),
         ]
     )
+
+
+@pytest.fixture
+def import_without_redis_installed(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], ModuleType]:
+    def _import(module_name: str) -> ModuleType:
+        package_name, _, attribute = module_name.rpartition(".")
+        monkeypatch.setattr(sys.modules[package_name], attribute, sys.modules[module_name])
+        monkeypatch.setitem(sys.modules, "redis", None)
+        for name in tuple(sys.modules):
+            if name.startswith("redis."):
+                monkeypatch.setitem(sys.modules, name, None)
+        monkeypatch.delitem(sys.modules, module_name)
+        return importlib.import_module(module_name)
+
+    return _import

tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py

@@ -1,10 +1,13 @@
 import json
 import logging
+from collections.abc import Callable
 from datetime import datetime, timezone
+from types import ModuleType
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from redis.exceptions import RedisError
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
@@ -12,6 +15,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import DailyTagSpendTransaction
 from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
 from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
+from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue
 from litellm.proxy.proxy_server import ProxyStartupEvent
 
 
@@ -572,8 +576,6 @@ async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer
 @pytest.mark.asyncio
 async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache):
     """A Redis failure during restore must not propagate to the caller's finally block."""
-    from redis.exceptions import RedisError
-
     mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down"))
 
     await redis_update_buffer.restore_transactions_to_redis(
@@ -909,11 +911,10 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_
 
 
 @pytest.mark.asyncio
-async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails():
-    from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue
-
+@pytest.mark.parametrize("failure", [ConnectionError("redis down"), RedisError("redis down")])
+async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(failure: Exception):
     redis_cache = MagicMock()
-    redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down"))
+    redis_cache.async_rpush_pipeline = AsyncMock(side_effect=failure)
     queue = ToolDiscoveryQueue()
     queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
     queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
@@ -921,3 +922,18 @@ async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails()
     await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue)
 
     assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}]
+
+
+@pytest.mark.asyncio
+async def test_buffer_without_redis_installed_skips_redis_work_when_no_cache_is_configured(
+    import_without_redis_installed: Callable[[str], ModuleType],
+):
+    module: Final = import_without_redis_installed("litellm.proxy.db.db_transaction_queue.redis_update_buffer")
+    buffer: Final[RedisUpdateBuffer] = module.RedisUpdateBuffer(redis_cache=None)
+    queue: Final = ToolDiscoveryQueue()
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+
+    await buffer.restore_transactions_to_redis(db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}})
+    await buffer.store_tool_discoveries_in_redis(queue)
+
+    assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]

tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py

@@ -6,8 +6,9 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of
 
 import asyncio
 import time
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from datetime import datetime, timedelta
+from types import ModuleType
 from typing import Final
 from unittest.mock import AsyncMock, patch
 
@@ -3677,3 +3678,16 @@ async def test_a_worker_without_redis_reads_the_models_saturation_from_its_own_m
     group: Final = handler.llm_router.registry().group("shared-model")
     assert group is not None
     assert await handler._check_model_saturation("shared-model", group) == pytest.approx(0.7)
+
+
+@pytest.mark.asyncio
+async def test_the_limiter_loads_and_admits_requests_without_redis_installed(
+    monkeypatch: pytest.MonkeyPatch, import_without_redis_installed: Callable[[str], ModuleType]
+):
+    module: Final = import_without_redis_installed("litellm.proxy.hooks.dynamic_rate_limiter_v3")
+    handler: Final[PROXY_DynamicRateLimitHandlerV3] = module.PROXY_DynamicRateLimitHandlerV3(
+        internal_usage_cache=DualCache()
+    )
+    handler.update_variables(llm_router=_borrowing_handler(monkeypatch, borrows=False).llm_router)
+
+    await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")