diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 459f95972ce..ca9665b26d7 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/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 diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 65da5630f33..eb79dd61526 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/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(