phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(caching.py): support default ttl for caching

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-25 13:40:17 -07:00
parent e1c75c009f
commit 591a0a376e
2 changed files with 23 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/caching.py
View file

@ -874,6 +874,7 @@ class Cache:
port: Optional[str] = None, port: Optional[str] = None,
password: Optional[str] = None, password: Optional[str] = None,
namespace: Optional[str] = None, namespace: Optional[str] = None,
ttl: Optional[float] = None,
similarity_threshold: Optional[float] = None, similarity_threshold: Optional[float] = None,
supported_call_types: Optional[ supported_call_types: Optional[
List[ List[
@ -967,6 +968,7 @@ class Cache:
self.supported_call_types = supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] self.supported_call_types = supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"]
self.type = type self.type = type
self.namespace = namespace self.namespace = namespace
self.ttl = ttl
def get_cache_key(self, *args, **kwargs): def get_cache_key(self, *args, **kwargs):
""" """
@ -1206,6 +1208,9 @@ class Cache:
if isinstance(result, OpenAIObject): if isinstance(result, OpenAIObject):
result = result.model_dump_json() result = result.model_dump_json()
## DEFAULT TTL ##
if self.ttl is not None:
kwargs["ttl"] = self.ttl
## Get Cache-Controls ## ## Get Cache-Controls ##
if kwargs.get("cache", None) is not None and isinstance( if kwargs.get("cache", None) is not None and isinstance(
kwargs.get("cache"), dict kwargs.get("cache"), dict
@ -1213,6 +1218,7 @@ class Cache:
for k, v in kwargs.get("cache").items(): for k, v in kwargs.get("cache").items():
if k == "ttl": if k == "ttl":
kwargs["ttl"] = v kwargs["ttl"] = v
cached_data = {"timestamp": time.time(), "response": result} cached_data = {"timestamp": time.time(), "response": result}
return cache_key, cached_data, kwargs return cache_key, cached_data, kwargs
else: else:

17
litellm/tests/test_caching.py
View file

@ -116,6 +116,23 @@ def test_caching_with_ttl():
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
def test_caching_with_default_ttl():
try:
litellm.set_verbose = True
litellm.cache = Cache(ttl=0)
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
litellm.success_callback = []
litellm._async_success_callback = []
assert response2["id"] != response1["id"]
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
def test_caching_with_cache_controls(): def test_caching_with_cache_controls():
try: try:
litellm.set_verbose = True litellm.set_verbose = True