mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
"""
|
|
Redis Cluster Cache implementation
|
|
|
|
Key differences:
|
|
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
|
|
"""
|
|
|
|
from typing import TYPE_CHECKING, Any, List, Optional, Union
|
|
|
|
from litellm.caching.redis_cache import RedisCache
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
from redis.asyncio import Redis, RedisCluster
|
|
from redis.asyncio.client import Pipeline
|
|
|
|
pipeline = Pipeline
|
|
async_redis_client = Redis
|
|
Span = Union[_Span, Any]
|
|
else:
|
|
pipeline = Any
|
|
async_redis_client = Any
|
|
Span = Any
|
|
|
|
|
|
class RedisClusterCache(RedisCache):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.redis_async_redis_cluster_client: Optional[RedisCluster] = None
|
|
self.redis_sync_redis_cluster_client: Optional[RedisCluster] = None
|
|
|
|
def init_async_client(self):
|
|
from redis.asyncio import RedisCluster
|
|
|
|
from .._redis import get_redis_async_client
|
|
|
|
if self.redis_async_redis_cluster_client:
|
|
return self.redis_async_redis_cluster_client
|
|
|
|
_redis_client = get_redis_async_client(
|
|
connection_pool=self.async_redis_conn_pool, **self.redis_kwargs
|
|
)
|
|
if isinstance(_redis_client, RedisCluster):
|
|
self.redis_async_redis_cluster_client = _redis_client
|
|
|
|
return _redis_client
|
|
|
|
def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
|
|
"""
|
|
Overrides `_run_redis_mget_operation` in redis_cache.py
|
|
"""
|
|
return self.redis_client.mget_nonatomic(keys=keys) # type: ignore
|
|
|
|
async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
|
|
"""
|
|
Overrides `_async_run_redis_mget_operation` in redis_cache.py
|
|
"""
|
|
async_redis_cluster_client = self.init_async_client()
|
|
return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore
|