diff --git a/litellm/_redis.py b/litellm/_redis.py index e6ac323ff5..a86ebd9ea9 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -78,6 +78,7 @@ def _get_redis_cluster_kwargs(client=None): available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection available_args.append("gcp_service_account") available_args.append("gcp_ssl_ca_certs") + available_args.append("max_connections") return available_args @@ -376,7 +377,7 @@ def get_redis_client(**env_overrides): def get_redis_async_client( - **env_overrides, + connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **env_overrides, ) -> Union[async_redis.Redis, async_redis.RedisCluster]: redis_kwargs = _get_redis_client_logic(**env_overrides) if "url" in redis_kwargs and redis_kwargs["url"] is not None: @@ -447,6 +448,10 @@ def get_redis_async_client( if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs: return _init_async_redis_sentinel(redis_kwargs) _pretty_print_redis_config(redis_kwargs=redis_kwargs) + + if connection_pool is not None: + redis_kwargs["connection_pool"] = connection_pool + return async_redis.Redis( **redis_kwargs, ) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 3ea1811549..4709faea4b 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,6 +1,8 @@ -from litellm._redis import get_redis_url_from_environment +from litellm._redis import get_redis_url_from_environment, _get_redis_cluster_kwargs, get_redis_async_client import os import pytest +from unittest.mock import MagicMock, patch +import redis.asyncio as async_redis def test_get_redis_url_from_environment_single_url(monkeypatch): """Test when REDIS_URL is directly provided""" @@ -117,3 +119,51 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): # Check the error message assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) + +def test_max_connections_in_cluster_kwargs(): + """Test that max_connections is included in Redis cluster kwargs""" + kwargs = _get_redis_cluster_kwargs() + assert "max_connections" in kwargs, "max_connections should be in available Redis cluster kwargs" + +def test_get_redis_async_client_with_connection_pool(): + """Test that connection_pool parameter is properly passed to Redis client""" + # Create a mock connection pool + mock_pool = MagicMock(spec=async_redis.BlockingConnectionPool) + + # Mock the Redis client creation + with patch('litellm._redis.async_redis.Redis') as mock_redis, \ + patch('litellm._redis._get_redis_client_logic') as mock_logic: + + # Configure mock to return basic redis kwargs + mock_logic.return_value = { + "host": "localhost", + "port": 6379, + "db": 0 + } + + # Call get_redis_async_client with connection_pool + get_redis_async_client(connection_pool=mock_pool) + + # Verify Redis was called with connection_pool in kwargs + call_kwargs = mock_redis.call_args[1] + assert "connection_pool" in call_kwargs, "connection_pool should be passed to Redis client" + assert call_kwargs["connection_pool"] == mock_pool, "connection_pool should match the provided pool" + +def test_get_redis_async_client_without_connection_pool(): + """Test that Redis client works without connection_pool parameter""" + with patch('litellm._redis.async_redis.Redis') as mock_redis, \ + patch('litellm._redis._get_redis_client_logic') as mock_logic: + + # Configure mock to return basic redis kwargs + mock_logic.return_value = { + "host": "localhost", + "port": 6379, + "db": 0 + } + + # Call get_redis_async_client without connection_pool + get_redis_async_client() + + # Verify Redis was called without connection_pool in kwargs + call_kwargs = mock_redis.call_args[1] + assert "connection_pool" not in call_kwargs, "connection_pool should not be in kwargs when not provided"