mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +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
39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
"""
|
|
Add the event loop to the cache key, to prevent event loop closed errors.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from .in_memory_cache import InMemoryCache
|
|
|
|
|
|
class LLMClientCache(InMemoryCache):
|
|
def update_cache_key_with_event_loop(self, key):
|
|
"""
|
|
Add the event loop to the cache key, to prevent event loop closed errors.
|
|
If none, use the key as is.
|
|
"""
|
|
try:
|
|
event_loop = asyncio.get_event_loop()
|
|
stringified_event_loop = str(id(event_loop))
|
|
return f"{key}-{stringified_event_loop}"
|
|
except Exception: # handle no current event loop
|
|
return key
|
|
|
|
def set_cache(self, key, value, **kwargs):
|
|
key = self.update_cache_key_with_event_loop(key)
|
|
return super().set_cache(key, value, **kwargs)
|
|
|
|
async def async_set_cache(self, key, value, **kwargs):
|
|
key = self.update_cache_key_with_event_loop(key)
|
|
return await super().async_set_cache(key, value, **kwargs)
|
|
|
|
def get_cache(self, key, **kwargs):
|
|
key = self.update_cache_key_with_event_loop(key)
|
|
|
|
return super().get_cache(key, **kwargs)
|
|
|
|
async def async_get_cache(self, key, **kwargs):
|
|
key = self.update_cache_key_with_event_loop(key)
|
|
|
|
return await super().async_get_cache(key, **kwargs)
|