Merge 792394c8b9 into 40be7fa1a8

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
@@ -542,6 +541,8 @@ class RedisUpdateBuffer:
         if self.redis_cache is None:
             return
 
+        from redis.exceptions import RedisError
+
         restore_configs: Final = (
             (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY),
             (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY),
@@ -584,6 +585,9 @@ 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

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

@@ -1,6 +1,9 @@
+import builtins
+import importlib.util
 import json
 import logging
 from datetime import datetime, timezone
+from pathlib import Path
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -15,6 +18,30 @@ from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdat
 from litellm.proxy.proxy_server import ProxyStartupEvent
 
 
+def test_importing_module_does_not_require_redis(monkeypatch: pytest.MonkeyPatch) -> None:
+    original_import: Final = builtins.__import__
+
+    def import_without_redis(
+        name: str,
+        globals: dict[str, object] | None = None,
+        locals: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "redis" or name.startswith("redis."):
+            raise ModuleNotFoundError("No module named 'redis'")
+        return original_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", import_without_redis)
+    module_path: Final = Path(__file__).parents[5] / "litellm/proxy/db/db_transaction_queue/redis_update_buffer.py"
+    spec: Final = importlib.util.spec_from_file_location("redis_update_buffer_without_redis", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+
+    module: Final = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+
 def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]:
     return {
         "research_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions": DailyTagSpendTransaction(