fix(_redis.py): fix args passed to redis.from_url argument

This commit is contained in:
Krrish Dholakia 2024-04-19 15:27:11 -07:00
parent 08cf77623b
commit 84685b5f34
2 changed files with 41 additions and 8 deletions

View file

@ -32,6 +32,25 @@ def _get_redis_kwargs():
return available_args return available_args
def _get_redis_url_kwargs(client=None):
if client is None:
client = redis.Redis.from_url
arg_spec = inspect.getfullargspec(redis.Redis.from_url)
# Only allow primitive arguments
exclude_args = {
"self",
"connection_pool",
"retry",
}
include_args = ["url"]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
return available_args
def _get_redis_env_kwarg_mapping(): def _get_redis_env_kwarg_mapping():
PREFIX = "REDIS_" PREFIX = "REDIS_"
@ -98,20 +117,31 @@ def _get_redis_client_logic(**env_overrides):
def get_redis_client(**env_overrides): def get_redis_client(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides) redis_kwargs = _get_redis_client_logic(**env_overrides)
if "url" in redis_kwargs and redis_kwargs["url"] is not None: if "url" in redis_kwargs and redis_kwargs["url"] is not None:
redis_kwargs.pop( args = _get_redis_url_kwargs()
"connection_pool", None url_kwargs = {}
) # redis.from_url doesn't support setting your own connection pool for arg in redis_kwargs:
return redis.Redis.from_url(**redis_kwargs) if arg in args:
url_kwargs[arg] = redis_kwargs[arg]
return redis.Redis.from_url(**url_kwargs)
return redis.Redis(**redis_kwargs) return redis.Redis(**redis_kwargs)
def get_redis_async_client(**env_overrides): def get_redis_async_client(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides) redis_kwargs = _get_redis_client_logic(**env_overrides)
if "url" in redis_kwargs and redis_kwargs["url"] is not None: if "url" in redis_kwargs and redis_kwargs["url"] is not None:
redis_kwargs.pop( args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
"connection_pool", None url_kwargs = {}
) # redis.from_url doesn't support setting your own connection pool for arg in redis_kwargs:
return async_redis.Redis.from_url(**redis_kwargs) if arg in args:
url_kwargs[arg] = redis_kwargs[arg]
else:
litellm.print_verbose(
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(
arg
)
)
return async_redis.Redis.from_url(**url_kwargs)
return async_redis.Redis( return async_redis.Redis(
socket_timeout=5, socket_timeout=5,
**redis_kwargs, **redis_kwargs,

View file

@ -174,6 +174,9 @@ class RedisCache(BaseCache):
except Exception: except Exception:
pass pass
### SYNC HEALTH PING ###
self.redis_client.ping()
def init_async_client(self): def init_async_client(self):
from ._redis import get_redis_async_client from ._redis import get_redis_async_client