litellm/tests/local_testing/test_caching.py

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

2825 lines
89 KiB
Python
Raw Normal View History

import os
import sys
2023-10-30 20:35:55 +08:00
import time
2023-08-17 12:44:50 +08:00
import traceback
from litellm._uuid import uuid
2023-08-17 12:44:50 +08:00
from dotenv import load_dotenv
2023-08-19 02:05:05 +08:00
2023-08-17 12:44:50 +08:00
load_dotenv()
import os
import json
2023-08-19 02:05:05 +08:00
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import hashlib
import random
2023-08-17 12:44:50 +08:00
import pytest
2023-08-17 12:44:50 +08:00
import litellm
from litellm import aembedding, completion, embedding
from litellm.caching.caching import Cache
from redis.asyncio import RedisCluster
from litellm.caching.redis_cluster_cache import RedisClusterCache
from unittest.mock import AsyncMock, patch, MagicMock, call
import datetime
from datetime import timedelta
2023-11-01 09:32:31 +08:00
# litellm.set_verbose=True
2023-08-17 12:44:50 +08:00
2023-08-18 01:41:14 +08:00
messages = [{"role": "user", "content": "who is ishaan Github? "}]
2023-08-29 04:14:58 +08:00
# comment
2023-08-29 00:20:50 +08:00
2023-12-10 05:56:53 +08:00
import random
import string
2023-12-25 16:40:38 +08:00
2023-12-10 05:56:53 +08:00
def generate_random_word(length=4):
letters = string.ascii_lowercase
2023-12-25 16:40:38 +08:00
return "".join(random.choice(letters) for _ in range(length))
2023-12-10 05:56:53 +08:00
2023-08-29 04:11:54 +08:00
messages = [{"role": "user", "content": "who is ishaan 5222"}]
2023-12-25 16:40:38 +08:00
@pytest.mark.asyncio
async def test_dual_cache_async_batch_get_cache():
"""
Unit testing for Dual Cache async_batch_get_cache()
- 2 item query
- in_memory result has a partial hit (1/2)
- hit redis for the other -> expect to return None
- expect result = [in_memory_result, None]
"""
from litellm.caching.caching import DualCache, InMemoryCache, RedisCache
in_memory_cache = InMemoryCache()
redis_cache = RedisCache() # get credentials from environment
dual_cache = DualCache(in_memory_cache=in_memory_cache, redis_cache=redis_cache)
with patch.object(
dual_cache.redis_cache, "async_batch_get_cache", new=AsyncMock()
) as mock_redis_cache:
mock_redis_cache.return_value = {"test_value_2": None, "test_value": "hello"}
await dual_cache.async_batch_get_cache(keys=["test_value", "test_value_2"])
await dual_cache.async_batch_get_cache(keys=["test_value", "test_value_2"])
assert mock_redis_cache.call_count == 1
def test_dual_cache_batch_get_cache():
"""
Unit testing for Dual Cache batch_get_cache()
- 2 item query
- in_memory result has a partial hit (1/2)
- hit redis for the other -> expect to return None
- expect result = [in_memory_result, None]
"""
from litellm.caching.caching import DualCache, InMemoryCache, RedisCache
in_memory_cache = InMemoryCache()
redis_cache = RedisCache() # get credentials from environment
dual_cache = DualCache(in_memory_cache=in_memory_cache, redis_cache=redis_cache)
in_memory_cache.set_cache(key="test_value", value="hello world")
result = dual_cache.batch_get_cache(
keys=["test_value", "test_value_2"], parent_otel_span=None
)
assert result[0] == "hello world"
assert result[1] == None
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_batch_get_cache_with_none_keys(sync_mode):
"""
Unit testing for RedisCache batch_get_cache() and async_batch_get_cache()
- test with None keys. Ensure it can safely handle when keys are None.
- expect result = {key: None}
"""
from litellm.caching.caching import RedisCache
litellm._turn_on_debug()
redis_cache = RedisCache(
host=os.environ.get("REDIS_HOST"),
port=os.environ.get("REDIS_PORT"),
password=os.environ.get("REDIS_PASSWORD"),
)
keys_to_lookup = [
None,
f"test_value_{uuid.uuid4()}",
None,
f"test_value_2_{uuid.uuid4()}",
None,
f"test_value_3_{uuid.uuid4()}",
]
if sync_mode:
result = redis_cache.batch_get_cache(key_list=keys_to_lookup)
print("result from batch_get_cache=", result)
else:
result = await redis_cache.async_batch_get_cache(key_list=keys_to_lookup)
print("result from async_batch_get_cache=", result)
expected_result = {}
for key in keys_to_lookup:
if key is None:
continue
expected_result[key] = None
assert result == expected_result
# @pytest.mark.skip(reason="")
def test_caching_dynamic_args(): # test in memory cache
try:
litellm.set_verbose = True
_redis_host_env = os.environ.pop("REDIS_HOST")
_redis_port_env = os.environ.pop("REDIS_PORT")
_redis_password_env = os.environ.pop("REDIS_PASSWORD")
litellm.cache = Cache(
type="redis",
host=_redis_host_env,
port=_redis_port_env,
password=_redis_password_env,
)
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 = []
if (
response2["choices"][0]["message"]["content"]
!= response1["choices"][0]["message"]["content"]
):
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
os.environ["REDIS_HOST"] = _redis_host_env
os.environ["REDIS_PORT"] = _redis_port_env
os.environ["REDIS_PASSWORD"] = _redis_password_env
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
2023-12-25 16:40:38 +08:00
def test_caching_v2(): # test in memory cache
2023-08-29 04:11:54 +08:00
try:
2023-12-25 16:40:38 +08:00
litellm.set_verbose = True
2023-08-29 04:11:54 +08:00
litellm.cache = Cache()
2023-09-09 11:15:14 +08:00
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
2023-08-29 04:11:54 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
2023-12-25 16:40:38 +08:00
litellm.cache = None # disable cache
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-25 16:40:38 +08:00
if (
response2["choices"][0]["message"]["content"]
!= response1["choices"][0]["message"]["content"]
):
2023-08-29 04:11:54 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
2023-08-29 04:11:54 +08:00
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
2023-09-09 11:15:14 +08:00
2023-12-25 16:40:38 +08:00
# test_caching_v2()
2023-08-29 04:11:54 +08:00
def test_caching_with_ttl():
try:
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0
)
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["choices"][0]["message"]["content"]
!= response1["choices"][0]["message"]["content"]
)
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
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}")
@pytest.mark.parametrize(
"sync_flag",
[True, False],
)
@pytest.mark.asyncio
async def test_caching_with_cache_controls(sync_flag):
try:
litellm.set_verbose = True
litellm.cache = Cache()
message = [{"role": "user", "content": f"Hey, how's it going? {uuid.uuid4()}"}]
if sync_flag:
## TTL = 0
response1 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"ttl": 0}
)
response2 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 10}
)
assert response2["id"] != response1["id"]
else:
## TTL = 0
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
cache={"ttl": 0},
mock_response="Hello world",
)
await asyncio.sleep(10)
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
cache={"s-maxage": 10},
mock_response="Hello world",
)
assert response2["id"] != response1["id"]
message = [{"role": "user", "content": f"Hey, how's it going? {uuid.uuid4()}"}]
## TTL = 5
if sync_flag:
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
cache={"ttl": 5},
mock_response="Hello world",
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
cache={"s-maxage": 5},
mock_response="Hello world",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response2["id"] == response1["id"]
else:
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
cache={"ttl": 25},
mock_response="Hello world",
)
await asyncio.sleep(10)
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
cache={"s-maxage": 25},
mock_response="Hello world",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response2["id"] == response1["id"]
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
# test_caching_with_cache_controls()
2025-08-31 04:21:14 +08:00
@pytest.mark.flaky(retries=3, delay=1)
2023-08-29 05:07:51 +08:00
def test_caching_with_models_v2():
2023-12-25 16:40:38 +08:00
messages = [
{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}
]
2023-08-29 04:11:54 +08:00
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
2023-09-09 11:15:14 +08:00
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
2025-09-28 01:08:32 +08:00
response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True)
2023-08-29 04:11:54 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-25 16:40:38 +08:00
if (
response3["choices"][0]["message"]["content"]
== response2["choices"][0]["message"]["content"]
):
2023-08-29 04:11:54 +08:00
# if models are different, it should not return cached response
print(f"response2: {response2}")
print(f"response3: {response3}")
pytest.fail(f"Error occurred:")
2023-12-25 16:40:38 +08:00
if (
response1["choices"][0]["message"]["content"]
!= response2["choices"][0]["message"]["content"]
):
2023-08-29 04:11:54 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
2023-12-25 16:40:38 +08:00
2023-10-21 06:00:40 +08:00
# test_caching_with_models_v2()
2023-08-29 05:07:51 +08:00
2024-08-06 02:23:49 +08:00
def c():
2024-08-06 02:23:20 +08:00
litellm.enable_caching_on_provider_specific_optional_params = True
messages = [
{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}
]
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
top_k=10,
caching=True,
mock_response="Hello: {}".format(uuid.uuid4()),
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
top_k=10,
caching=True,
mock_response="Hello: {}".format(uuid.uuid4()),
)
response3 = completion(
model="gpt-3.5-turbo",
messages=messages,
top_k=9,
caching=True,
mock_response="Hello: {}".format(uuid.uuid4()),
)
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
if (
response3["choices"][0]["message"]["content"]
== response2["choices"][0]["message"]["content"]
):
# if models are different, it should not return cached response
print(f"response2: {response2}")
print(f"response3: {response3}")
pytest.fail(f"Error occurred:")
if (
response1["choices"][0]["message"]["content"]
!= response2["choices"][0]["message"]["content"]
):
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
2024-08-06 02:23:20 +08:00
litellm.enable_caching_on_provider_specific_optional_params = False
2023-12-25 16:40:38 +08:00
embedding_large_text = (
"""
2023-08-29 05:07:51 +08:00
small text
2023-12-25 16:40:38 +08:00
"""
* 5
)
2023-08-29 05:07:51 +08:00
2023-08-29 12:56:57 +08:00
# # test_caching_with_models()
2023-08-29 05:07:51 +08:00
def test_embedding_caching():
import time
2023-12-25 16:40:38 +08:00
# litellm.set_verbose = True
2023-08-29 05:07:51 +08:00
litellm.cache = Cache()
text_to_embed = [embedding_large_text]
start_time = time.time()
2023-12-25 16:40:38 +08:00
embedding1 = embedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
2023-08-29 05:07:51 +08:00
end_time = time.time()
print(f"Embedding 1 response time: {end_time - start_time} seconds")
2023-08-29 12:56:57 +08:00
2023-08-29 05:07:51 +08:00
time.sleep(1)
start_time = time.time()
2023-12-25 16:40:38 +08:00
embedding2 = embedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
2023-08-29 05:07:51 +08:00
end_time = time.time()
# print(f"embedding2: {embedding2}")
2023-08-29 05:07:51 +08:00
print(f"Embedding 2 response time: {end_time - start_time} seconds")
2023-10-30 20:35:55 +08:00
2023-08-29 05:07:51 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-25 16:40:38 +08:00
assert end_time - start_time <= 0.1 # ensure 2nd response comes in in under 0.1 s
if embedding2["data"][0]["embedding"] != embedding1["data"][0]["embedding"]:
2023-08-29 05:07:51 +08:00
print(f"embedding1: {embedding1}")
print(f"embedding2: {embedding2}")
pytest.fail("Error occurred: Embedding caching failed")
2023-12-25 16:40:38 +08:00
2023-12-05 12:50:06 +08:00
# test_embedding_caching()
2023-10-14 12:16:41 +08:00
@pytest.mark.asyncio
async def test_embedding_caching_individual_items_and_then_list():
litellm._turn_on_debug()
litellm.cache = Cache()
text_to_embed = [
"hello",
"world",
]
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[0], caching=True
)
initial_prompt_tokens = embedding1.usage.prompt_tokens
await asyncio.sleep(1)
embedding2 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[1], caching=True
)
await asyncio.sleep(1)
embedding3 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
final_prompt_tokens = embedding3.usage.prompt_tokens
assert embedding3["data"][0]["embedding"] == embedding1["data"][0]["embedding"]
assert embedding3["data"][1]["embedding"] == embedding2["data"][0]["embedding"]
assert embedding3._hidden_params["cache_hit"] == True
assert embedding3.usage.prompt_tokens != 0
## with new input, check that prompt tokens increase
additional_text = "this is a new text"
text_to_embed.append(additional_text)
embedding4 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens
@pytest.mark.asyncio
async def test_embedding_caching_individual_items():
litellm.cache = Cache()
text_to_embed = "hello"
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
await asyncio.sleep(1)
embedding3 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True
)
final_prompt_tokens = embedding3.usage.prompt_tokens
assert embedding3["data"][0]["embedding"] == embedding1["data"][0]["embedding"]
assert len(embedding3.data) == 1
assert embedding3._hidden_params["cache_hit"] == True
assert embedding3.usage.prompt_tokens != 0
2023-10-14 12:16:41 +08:00
def test_embedding_caching_azure():
print("Testing azure embedding caching")
import time
2023-12-25 16:40:38 +08:00
2023-10-14 12:16:41 +08:00
litellm.cache = Cache()
text_to_embed = [embedding_large_text]
2023-12-25 16:40:38 +08:00
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_API_BASE"]
api_version = os.environ["AZURE_API_VERSION"]
2023-10-14 12:16:41 +08:00
2023-12-25 16:40:38 +08:00
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_KEY"] = ""
2023-10-14 12:16:41 +08:00
start_time = time.time()
2023-10-14 13:18:52 +08:00
print("AZURE CONFIGS")
print(api_version)
print(api_key)
print(api_base)
2023-10-14 12:16:41 +08:00
embedding1 = embedding(
2025-09-28 03:34:49 +08:00
model="azure/text-embedding-ada-002",
2023-10-14 12:16:41 +08:00
input=["good morning from litellm", "this is another item"],
api_key=api_key,
api_base=api_base,
api_version=api_version,
2023-12-25 16:40:38 +08:00
caching=True,
2023-10-14 12:16:41 +08:00
)
end_time = time.time()
print(f"Embedding 1 response time: {end_time - start_time} seconds")
time.sleep(1)
start_time = time.time()
embedding2 = embedding(
2025-09-28 03:34:49 +08:00
model="azure/text-embedding-ada-002",
2023-10-14 12:16:41 +08:00
input=["good morning from litellm", "this is another item"],
api_key=api_key,
api_base=api_base,
api_version=api_version,
2023-12-25 16:40:38 +08:00
caching=True,
2023-10-14 12:16:41 +08:00
)
end_time = time.time()
print(f"Embedding 2 response time: {end_time - start_time} seconds")
2023-10-30 20:35:55 +08:00
2023-10-14 12:16:41 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-25 16:40:38 +08:00
assert end_time - start_time <= 0.1 # ensure 2nd response comes in in under 0.1 s
if embedding2["data"][0]["embedding"] != embedding1["data"][0]["embedding"]:
2023-10-14 12:16:41 +08:00
print(f"embedding1: {embedding1}")
print(f"embedding2: {embedding2}")
pytest.fail("Error occurred: Embedding caching failed")
2023-10-30 20:35:55 +08:00
2023-12-25 16:40:38 +08:00
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_API_KEY"] = api_key
2023-10-14 12:16:41 +08:00
2023-10-25 03:06:24 +08:00
# test_embedding_caching_azure()
2023-08-29 04:11:54 +08:00
2024-01-11 19:00:05 +08:00
@pytest.mark.asyncio
async def test_embedding_caching_azure_individual_items():
"""
Tests caching for individual items in an embedding list
2024-01-13 00:16:41 +08:00
- Cache an item
- call aembedding(..) with the item + 1 unique item
- compare to a 2nd aembedding (...) with 2 unique items
2024-01-11 19:00:05 +08:00
```
embedding_1 = ["hey how's it going", "I'm doing well"]
embedding_val_1 = embedding(...)
embedding_2 = ["hey how's it going", "I'm fine"]
embedding_val_2 = embedding(...)
assert embedding_val_1[0]["id"] == embedding_val_2[0]["id"]
```
"""
litellm.cache = Cache()
common_msg = f"hey how's it going {uuid.uuid4()}"
2024-01-13 00:16:41 +08:00
common_msg_2 = f"hey how's it going {uuid.uuid4()}"
embedding_1 = [common_msg]
2024-01-13 00:16:41 +08:00
embedding_2 = [
common_msg,
f"I'm fine {uuid.uuid4()}",
]
2024-01-11 19:00:05 +08:00
embedding_val_1 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002", input=embedding_1, caching=True
2024-01-11 19:00:05 +08:00
)
embedding_val_2 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002", input=embedding_2, caching=True
2024-01-11 19:00:05 +08:00
)
print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}")
assert embedding_val_2._hidden_params["cache_hit"] == True
2024-01-13 00:16:41 +08:00
@pytest.mark.asyncio
async def test_embedding_caching_azure_individual_items_reordered():
"""
Tests caching for individual items in an embedding list
- Cache an item
- call aembedding(..) with the item + 1 unique item
- compare to a 2nd aembedding (...) with 2 unique items
```
embedding_1 = ["hey how's it going", "I'm doing well"]
embedding_val_1 = embedding(...)
embedding_2 = ["hey how's it going", "I'm fine"]
embedding_val_2 = embedding(...)
assert embedding_val_1[0]["id"] == embedding_val_2[0]["id"]
```
"""
litellm.set_verbose = True
litellm.cache = Cache()
common_msg = f"{uuid.uuid4()}"
common_msg_2 = f"hey how's it going {uuid.uuid4()}"
embedding_1 = [common_msg_2, common_msg]
embedding_2 = [
common_msg,
f"I'm fine {uuid.uuid4()}",
]
embedding_val_1 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002", input=embedding_1, caching=True
)
print("embedding val 1", embedding_val_1)
embedding_val_2 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002", input=embedding_2, caching=True
)
print("embedding val 2", embedding_val_2)
print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}")
assert embedding_val_2._hidden_params["cache_hit"] == True
assert embedding_val_2.data[0]["embedding"] == embedding_val_1.data[1]["embedding"]
assert embedding_val_2.data[0]["index"] != embedding_val_1.data[1]["index"]
assert embedding_val_2.data[0]["index"] == 0
assert embedding_val_1.data[1]["index"] == 1
2024-04-11 07:46:56 +08:00
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
2024-04-11 07:46:56 +08:00
async def test_embedding_caching_base_64():
""" """
2024-04-20 08:02:15 +08:00
litellm.set_verbose = True
2024-04-11 07:46:56 +08:00
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
)
from litellm._uuid import uuid
2024-04-11 07:46:56 +08:00
inputs = [
f"{uuid.uuid4()} hello this is ishaan",
f"{uuid.uuid4()} hello this is ishaan again",
]
embedding_val_1 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002",
2024-04-11 07:46:56 +08:00
input=inputs,
caching=True,
encoding_format="base64",
)
2024-04-20 08:02:15 +08:00
await asyncio.sleep(5)
print("\n\nCALL2\n\n")
2024-04-11 07:46:56 +08:00
embedding_val_2 = await aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002",
2024-04-11 07:46:56 +08:00
input=inputs,
caching=True,
encoding_format="base64",
)
assert embedding_val_2._hidden_params["cache_hit"] == True
print(embedding_val_2)
print(embedding_val_1)
assert embedding_val_2.data[0]["embedding"] == embedding_val_1.data[0]["embedding"]
assert embedding_val_2.data[1]["embedding"] == embedding_val_1.data[1]["embedding"]
@pytest.mark.asyncio
async def test_embedding_caching_redis_ttl():
"""
Test default_in_redis_ttl is used for embedding caching
issue: https://github.com/BerriAI/litellm/issues/6010
"""
litellm.set_verbose = True
# Create a mock for the pipeline
mock_pipeline = AsyncMock()
mock_set = AsyncMock()
mock_pipeline.__aenter__.return_value.set = mock_set
# Patch the Redis class to return our mock
with patch("redis.asyncio.Redis.pipeline", return_value=mock_pipeline):
# Simulate the context manager behavior for the pipeline
litellm.cache = Cache(
type="redis",
host="dummy_host",
password="dummy_password",
Litellm dev 10 29 2024 (#6502) * fix(core_helpers.py): return None, instead of raising kwargs is None error Closes https://github.com/BerriAI/litellm/issues/6500 * docs(cost_tracking.md): cleanup doc * fix(vertex_and_google_ai_studio.py): handle function call with no params passed in Closes https://github.com/BerriAI/litellm/issues/6495 * test(test_router_timeout.py): add test for router timeout + retry logic * test: update test to use module level values * (fix) Prometheus - Log Postgres DB latency, status on prometheus (#6484) * fix logging DB fails on prometheus * unit testing log to otel wrapper * unit testing for service logger + prometheus * use LATENCY buckets for service logging * fix service logging * docs clarify vertex vs gemini * (router_strategy/) ensure all async functions use async cache methods (#6489) * fix router strat * use async set / get cache in router_strategy * add coverage for router strategy * fix imports * fix batch_get_cache * use async methods for least busy * fix least busy use async methods * fix test_dual_cache_increment * test async_get_available_deployment when routing_strategy="least-busy" * (fix) proxy - fix when `STORE_MODEL_IN_DB` should be set (#6492) * set store_model_in_db at the top * correctly use store_model_in_db global * (fix) `PrometheusServicesLogger` `_get_metric` should return metric in Registry (#6486) * fix logging DB fails on prometheus * unit testing log to otel wrapper * unit testing for service logger + prometheus * use LATENCY buckets for service logging * fix service logging * fix _get_metric in prom services logger * add clear doc string * unit testing for prom service logger * bump: version 1.51.0 → 1.51.1 * Add `azure/gpt-4o-mini-2024-07-18` to model_prices_and_context_window.json (#6477) * Update utils.py (#6468) Fixed missing keys * (perf) Litellm redis router fix - ~100ms improvement (#6483) * docs(exception_mapping.md): add missing exception types Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183 * fix(main.py): register custom model pricing with specific key Ensure custom model pricing is registered to the specific model+provider key combination * test: make testing more robust for custom pricing * fix(redis_cache.py): instrument otel logging for sync redis calls ensures complete coverage for all redis cache calls * refactor: pass parent_otel_span for redis caching calls in router allows for more observability into what calls are causing latency issues * test: update tests with new params * refactor: ensure e2e otel tracing for router * refactor(router.py): add more otel tracing acrosss router catch all latency issues for router requests * fix: fix linting error * fix(router.py): fix linting error * fix: fix test * test: fix tests * fix(dual_cache.py): pass ttl to redis cache * fix: fix param * perf(cooldown_cache.py): improve cooldown cache, to store cache results in memory for 5s, prevents redis call from being made on each request reduces 100ms latency per call with caching enabled on router * fix: fix test * fix(cooldown_cache.py): handle if a result is None * fix(cooldown_cache.py): add debug statements * refactor(dual_cache.py): move to using an in-memory check for batch get cache, to prevent redis from being hit for every call * fix(cooldown_cache.py): fix linting erropr * refactor(prometheus.py): move to using standard logging payload for reading the remaining request / tokens Ensures prometheus token tracking works for anthropic as well * fix: fix linting error * fix(redis_cache.py): make sure ttl is always int (handle float values) Fixes issue where redis_client.ex was not working correctly due to float ttl * fix: fix linting error * test: update test * fix: fix linting error --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Xingyao Wang <xingyao@all-hands.dev> Co-authored-by: vibhanshu-ob <115142120+vibhanshu-ob@users.noreply.github.com>
2024-10-30 13:04:16 +08:00
default_in_redis_ttl=2,
)
inputs = [
f"{uuid.uuid4()} hello this is ishaan",
f"{uuid.uuid4()} hello this is ishaan again",
]
# Call the embedding method
embedding_val_1 = await litellm.aembedding(
2025-09-28 01:08:32 +08:00
model="text-embedding-ada-002",
input=inputs,
encoding_format="base64",
)
await asyncio.sleep(3) # Wait for TTL to expire
# Check if set was called on the pipeline
mock_set.assert_called()
# Check if the TTL was set correctly
for call in mock_set.call_args_list:
args, kwargs = call
print(f"redis pipeline set args: {args}")
print(f"redis pipeline set kwargs: {kwargs}")
assert kwargs.get("ex") == datetime.timedelta(
Litellm dev 10 29 2024 (#6502) * fix(core_helpers.py): return None, instead of raising kwargs is None error Closes https://github.com/BerriAI/litellm/issues/6500 * docs(cost_tracking.md): cleanup doc * fix(vertex_and_google_ai_studio.py): handle function call with no params passed in Closes https://github.com/BerriAI/litellm/issues/6495 * test(test_router_timeout.py): add test for router timeout + retry logic * test: update test to use module level values * (fix) Prometheus - Log Postgres DB latency, status on prometheus (#6484) * fix logging DB fails on prometheus * unit testing log to otel wrapper * unit testing for service logger + prometheus * use LATENCY buckets for service logging * fix service logging * docs clarify vertex vs gemini * (router_strategy/) ensure all async functions use async cache methods (#6489) * fix router strat * use async set / get cache in router_strategy * add coverage for router strategy * fix imports * fix batch_get_cache * use async methods for least busy * fix least busy use async methods * fix test_dual_cache_increment * test async_get_available_deployment when routing_strategy="least-busy" * (fix) proxy - fix when `STORE_MODEL_IN_DB` should be set (#6492) * set store_model_in_db at the top * correctly use store_model_in_db global * (fix) `PrometheusServicesLogger` `_get_metric` should return metric in Registry (#6486) * fix logging DB fails on prometheus * unit testing log to otel wrapper * unit testing for service logger + prometheus * use LATENCY buckets for service logging * fix service logging * fix _get_metric in prom services logger * add clear doc string * unit testing for prom service logger * bump: version 1.51.0 → 1.51.1 * Add `azure/gpt-4o-mini-2024-07-18` to model_prices_and_context_window.json (#6477) * Update utils.py (#6468) Fixed missing keys * (perf) Litellm redis router fix - ~100ms improvement (#6483) * docs(exception_mapping.md): add missing exception types Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183 * fix(main.py): register custom model pricing with specific key Ensure custom model pricing is registered to the specific model+provider key combination * test: make testing more robust for custom pricing * fix(redis_cache.py): instrument otel logging for sync redis calls ensures complete coverage for all redis cache calls * refactor: pass parent_otel_span for redis caching calls in router allows for more observability into what calls are causing latency issues * test: update tests with new params * refactor: ensure e2e otel tracing for router * refactor(router.py): add more otel tracing acrosss router catch all latency issues for router requests * fix: fix linting error * fix(router.py): fix linting error * fix: fix test * test: fix tests * fix(dual_cache.py): pass ttl to redis cache * fix: fix param * perf(cooldown_cache.py): improve cooldown cache, to store cache results in memory for 5s, prevents redis call from being made on each request reduces 100ms latency per call with caching enabled on router * fix: fix test * fix(cooldown_cache.py): handle if a result is None * fix(cooldown_cache.py): add debug statements * refactor(dual_cache.py): move to using an in-memory check for batch get cache, to prevent redis from being hit for every call * fix(cooldown_cache.py): fix linting erropr * refactor(prometheus.py): move to using standard logging payload for reading the remaining request / tokens Ensures prometheus token tracking works for anthropic as well * fix: fix linting error * fix(redis_cache.py): make sure ttl is always int (handle float values) Fixes issue where redis_client.ex was not working correctly due to float ttl * fix: fix linting error * test: update test * fix: fix linting error --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Xingyao Wang <xingyao@all-hands.dev> Co-authored-by: vibhanshu-ob <115142120+vibhanshu-ob@users.noreply.github.com>
2024-10-30 13:04:16 +08:00
seconds=2
) # Check if TTL is set to 2.5 seconds
2024-01-13 00:16:41 +08:00
@pytest.mark.asyncio
async def test_redis_cache_basic():
"""
Init redis client
- write to client
- read from client
"""
litellm.set_verbose = False
random_number = random.randint(
1, 100000
) # add a random number to ensure it's always adding / reading from cache
messages = [
{"role": "user", "content": f"write a one sentence poem about: {random_number}"}
]
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
)
cache_key = litellm.cache.get_cache_key(
model="gpt-3.5-turbo",
messages=messages,
)
print(f"cache_key: {cache_key}")
litellm.cache.add_cache(result=response1, cache_key=cache_key)
print(f"cache key pre async get: {cache_key}")
stored_val = await litellm.cache.async_get_cache(
model="gpt-3.5-turbo",
messages=messages,
)
print(f"stored_val: {stored_val}")
assert stored_val["id"] == response1.id
2024-01-11 19:00:05 +08:00
2024-03-26 09:04:04 +08:00
@pytest.mark.asyncio
(perf) Litellm redis router fix - ~100ms improvement (#6483) * docs(exception_mapping.md): add missing exception types Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183 * fix(main.py): register custom model pricing with specific key Ensure custom model pricing is registered to the specific model+provider key combination * test: make testing more robust for custom pricing * fix(redis_cache.py): instrument otel logging for sync redis calls ensures complete coverage for all redis cache calls * refactor: pass parent_otel_span for redis caching calls in router allows for more observability into what calls are causing latency issues * test: update tests with new params * refactor: ensure e2e otel tracing for router * refactor(router.py): add more otel tracing acrosss router catch all latency issues for router requests * fix: fix linting error * fix(router.py): fix linting error * fix: fix test * test: fix tests * fix(dual_cache.py): pass ttl to redis cache * fix: fix param * perf(cooldown_cache.py): improve cooldown cache, to store cache results in memory for 5s, prevents redis call from being made on each request reduces 100ms latency per call with caching enabled on router * fix: fix test * fix(cooldown_cache.py): handle if a result is None * fix(cooldown_cache.py): add debug statements * refactor(dual_cache.py): move to using an in-memory check for batch get cache, to prevent redis from being hit for every call * fix(cooldown_cache.py): fix linting erropr
2024-10-30 04:58:29 +08:00
@pytest.mark.flaky(retries=3, delay=1)
2024-03-26 09:04:04 +08:00
async def test_redis_batch_cache_write():
"""
Init redis client
- write to client
- read from client
"""
litellm.set_verbose = True
from litellm._uuid import uuid
2024-03-26 09:04:04 +08:00
messages = [
{"role": "user", "content": f"write a one sentence poem about: {uuid.uuid4()}"},
]
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
redis_flush_size=2,
)
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
)
response2 = await litellm.acompletion(
model="anthropic/claude-3-opus-20240229",
messages=messages,
mock_response="good morning from this test",
)
# we hit the flush size, this will now send to redis
await asyncio.sleep(2)
response4 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
)
assert response1.id == response4.id
2023-08-29 13:10:15 +08:00
def test_redis_cache_completion():
litellm.set_verbose = False
2023-12-25 16:40:38 +08:00
random_number = random.randint(
1, 100000
) # add a random number to ensure it's always adding / reading from cache
messages = [
{"role": "user", "content": f"write a one sentence poem about: {random_number}"}
]
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
print("test2 for Redis Caching - non streaming")
2023-12-25 16:40:38 +08:00
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
max_tokens=20,
2023-12-25 16:40:38 +08:00
)
response2 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20
)
response3 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, temperature=0.5
)
2025-03-19 10:05:26 +08:00
response4 = completion(model="gpt-4o-mini", messages=messages, caching=True)
print("\nresponse 1", response1)
print("\nresponse 2", response2)
print("\nresponse 3", response3)
2023-11-24 10:17:11 +08:00
print("\nresponse 4", response4)
2023-08-29 13:10:15 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-11-24 10:17:11 +08:00
"""
1 & 2 should be exactly the same
1 & 3 should be different, since input params are diff
1 & 4 should be diff, since models are diff
"""
2023-12-25 16:40:38 +08:00
if (
response1["choices"][0]["message"]["content"]
!= response2["choices"][0]["message"]["content"]
): # 1 and 2 should be the same
2023-11-24 10:17:11 +08:00
# 1&2 have the exact same input params. This MUST Be a CACHE HIT
2023-08-29 13:10:15 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
2023-12-25 16:40:38 +08:00
if (
response1["choices"][0]["message"]["content"]
== response3["choices"][0]["message"]["content"]
):
2023-11-24 10:17:11 +08:00
# if input params like seed, max_tokens are diff it should NOT be a cache hit
print(f"response1: {response1}")
print(f"response3: {response3}")
2023-12-25 16:40:38 +08:00
pytest.fail(
f"Response 1 == response 3. Same model, diff params shoudl not cache Error occurred:"
)
if (
response1["choices"][0]["message"]["content"]
== response4["choices"][0]["message"]["content"]
):
2023-11-24 10:17:11 +08:00
# if models are different, it should not return cached response
print(f"response1: {response1}")
print(f"response4: {response4}")
pytest.fail(f"Error occurred:")
2023-08-29 13:10:15 +08:00
assert response1.id == response2.id
assert response1.created == response2.created
assert response1.choices[0].message.content == response2.choices[0].message.content
2023-12-05 12:50:06 +08:00
# test_redis_cache_completion()
2023-08-29 13:10:15 +08:00
2023-12-25 16:40:38 +08:00
2023-12-09 03:50:36 +08:00
def test_redis_cache_completion_stream():
try:
2023-12-09 11:10:35 +08:00
litellm.success_callback = []
litellm._async_success_callback = []
litellm.callbacks = []
2023-12-09 05:14:59 +08:00
litellm.set_verbose = True
2023-12-25 16:40:38 +08:00
random_number = random.randint(
1, 100000
) # add a random number to ensure it's always adding / reading from cache
messages = [
{
"role": "user",
"content": f"write a one sentence poem about: {random_number}",
}
]
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
2023-12-09 03:50:36 +08:00
print("test for caching, streaming + completion")
2023-12-25 16:40:38 +08:00
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=40,
temperature=0.2,
stream=True,
)
response_1_id = ""
2023-12-09 03:50:36 +08:00
for chunk in response1:
print(chunk)
response_1_id = chunk.id
2023-12-09 03:50:36 +08:00
time.sleep(0.5)
2023-12-25 16:40:38 +08:00
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=40,
temperature=0.2,
stream=True,
)
response_2_id = ""
2023-12-09 03:50:36 +08:00
for chunk in response2:
print(chunk)
response_2_id = chunk.id
2023-12-25 16:40:38 +08:00
assert (
response_1_id == response_2_id
), f"Response 1 != Response 2. Same params, Response 1{response_1_id} != Response 2{response_2_id}"
2023-12-09 11:10:35 +08:00
litellm.success_callback = []
2023-12-09 03:50:36 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-09 03:50:36 +08:00
except Exception as e:
print(e)
2023-12-09 11:10:35 +08:00
litellm.success_callback = []
2023-12-09 03:50:36 +08:00
raise e
"""
1 & 2 should be exactly the same
"""
2023-12-25 16:40:38 +08:00
2024-01-11 19:00:05 +08:00
# test_redis_cache_completion_stream()
2023-12-10 05:56:53 +08:00
2024-08-22 06:05:18 +08:00
@pytest.mark.skip(reason="Local test. Requires running redis cluster locally.")
@pytest.mark.asyncio
async def test_redis_cache_cluster_init_unit_test():
try:
from redis.asyncio import RedisCluster as AsyncRedisCluster
from redis.cluster import RedisCluster
from litellm.caching.caching import RedisCache
litellm.set_verbose = True
# List of startup nodes
startup_nodes = [
{"host": "127.0.0.1", "port": "7001"},
]
resp = RedisCache(startup_nodes=startup_nodes)
assert isinstance(resp.redis_client, RedisCluster)
assert isinstance(resp.init_async_client(), AsyncRedisCluster)
resp = litellm.Cache(type="redis", redis_startup_nodes=startup_nodes)
assert isinstance(resp.cache, RedisCache)
assert isinstance(resp.cache.redis_client, RedisCluster)
assert isinstance(resp.cache.init_async_client(), AsyncRedisCluster)
except Exception as e:
print(f"{str(e)}\n\n{traceback.format_exc()}")
raise e
@pytest.mark.asyncio
@pytest.mark.skip(reason="Local test. Requires running redis cluster locally.")
async def test_redis_cache_cluster_init_with_env_vars_unit_test():
try:
import json
from redis.asyncio import RedisCluster as AsyncRedisCluster
from redis.cluster import RedisCluster
from litellm.caching.caching import RedisCache
litellm.set_verbose = True
# List of startup nodes
startup_nodes = [
{"host": "127.0.0.1", "port": "7001"},
{"host": "127.0.0.1", "port": "7003"},
{"host": "127.0.0.1", "port": "7004"},
{"host": "127.0.0.1", "port": "7005"},
{"host": "127.0.0.1", "port": "7006"},
{"host": "127.0.0.1", "port": "7007"},
]
# set startup nodes in environment variables
os.environ["REDIS_CLUSTER_NODES"] = json.dumps(startup_nodes)
print("REDIS_CLUSTER_NODES", os.environ["REDIS_CLUSTER_NODES"])
# unser REDIS_HOST, REDIS_PORT, REDIS_PASSWORD
os.environ.pop("REDIS_HOST", None)
os.environ.pop("REDIS_PORT", None)
os.environ.pop("REDIS_PASSWORD", None)
resp = RedisCache()
print("response from redis cache", resp)
assert isinstance(resp.redis_client, RedisCluster)
assert isinstance(resp.init_async_client(), AsyncRedisCluster)
resp = litellm.Cache(type="redis")
assert isinstance(resp.cache, RedisCache)
assert isinstance(resp.cache.redis_client, RedisCluster)
assert isinstance(resp.cache.init_async_client(), AsyncRedisCluster)
except Exception as e:
print(f"{str(e)}\n\n{traceback.format_exc()}")
raise e
2024-03-16 09:09:25 +08:00
@pytest.mark.asyncio
async def test_redis_cache_acompletion_stream():
2023-12-10 05:56:53 +08:00
try:
2024-03-16 09:09:25 +08:00
litellm.set_verbose = True
2023-12-10 05:56:53 +08:00
random_word = generate_random_word()
2023-12-25 16:40:38 +08:00
messages = [
{
"role": "user",
"content": f"write a one sentence poem about: {random_word}",
}
]
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
2023-12-10 05:56:53 +08:00
print("test for caching, streaming + completion")
response_1_content = ""
response_2_content = ""
2024-03-16 09:09:25 +08:00
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response1:
response_1_content += chunk.choices[0].delta.content or ""
print(response_1_content)
await asyncio.sleep(0.5)
2023-12-10 05:56:53 +08:00
print("\n\n Response 1 content: ", response_1_content, "\n\n")
2024-03-16 09:09:25 +08:00
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response2:
response_2_content += chunk.choices[0].delta.content or ""
print(response_2_content)
2023-12-10 05:56:53 +08:00
print("\nresponse 1", response_1_content)
print("\nresponse 2", response_2_content)
2023-12-25 16:40:38 +08:00
assert (
response_1_content == response_2_content
), f"Response 1 != Response 2. Same params, Response 1{response_1_content} != Response 2{response_2_content}"
2023-12-10 05:56:53 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-10 05:56:53 +08:00
except Exception as e:
2024-03-16 09:09:25 +08:00
print(f"{str(e)}\n\n{traceback.format_exc()}")
2023-12-10 05:56:53 +08:00
raise e
2023-12-25 16:40:38 +08:00
# test_redis_cache_acompletion_stream()
2023-12-09 03:50:36 +08:00
2023-12-25 16:40:38 +08:00
2024-04-13 03:32:21 +08:00
@pytest.mark.asyncio
async def test_redis_cache_atext_completion():
try:
litellm.set_verbose = True
prompt = f"write a one sentence poem about: {uuid.uuid4()}"
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
supported_call_types=["atext_completion"],
)
print("test for caching, atext_completion")
response1 = await litellm.atext_completion(
model="gpt-3.5-turbo-instruct", prompt=prompt, max_tokens=40, temperature=1
)
await asyncio.sleep(0.5)
print("\n\n Response 1 content: ", response1, "\n\n")
response2 = await litellm.atext_completion(
model="gpt-3.5-turbo-instruct", prompt=prompt, max_tokens=40, temperature=1
)
print(response2)
2024-04-13 11:37:56 +08:00
assert response1.id == response2.id
2024-04-13 03:32:21 +08:00
except Exception as e:
print(f"{str(e)}\n\n{traceback.format_exc()}")
raise e
2024-03-16 09:09:25 +08:00
@pytest.mark.asyncio
async def test_redis_cache_acompletion_stream_bedrock():
2023-12-12 00:43:00 +08:00
import asyncio
2023-12-25 16:40:38 +08:00
2023-12-12 00:43:00 +08:00
try:
litellm.set_verbose = True
random_word = generate_random_word()
2023-12-25 16:40:38 +08:00
messages = [
{
"role": "user",
"content": f"write a one sentence poem about: {random_word}",
}
]
litellm.cache = Cache(type="redis")
2023-12-12 00:43:00 +08:00
print("test for caching, streaming + completion")
response_1_content = ""
response_2_content = ""
2024-03-16 09:09:25 +08:00
response1 = await litellm.acompletion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
2024-03-16 09:09:25 +08:00
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response1:
print(chunk)
response_1_content += chunk.choices[0].delta.content or ""
print(response_1_content)
await asyncio.sleep(1)
2023-12-12 00:43:00 +08:00
print("\n\n Response 1 content: ", response_1_content, "\n\n")
2024-03-16 09:09:25 +08:00
response2 = await litellm.acompletion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
2024-03-16 09:09:25 +08:00
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response2:
print(chunk)
response_2_content += chunk.choices[0].delta.content or ""
print(response_2_content)
print("\nfinal response 1", response_1_content)
print("\nfinal response 2", response_2_content)
2023-12-25 16:40:38 +08:00
assert (
response_1_content == response_2_content
), f"Response 1 != Response 2. Same params, Response 1{response_1_content} != Response 2{response_2_content}"
2023-12-12 00:43:00 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-12 00:43:00 +08:00
except Exception as e:
print(e)
raise e
2023-12-25 16:40:38 +08:00
# @pytest.mark.skip(reason="AWS Suspended Account")
@pytest.mark.parametrize("sync_mode", [True, False])
2024-03-16 09:09:25 +08:00
@pytest.mark.asyncio
async def test_s3_cache_stream_azure(sync_mode):
2024-01-03 17:43:39 +08:00
try:
litellm.set_verbose = True
random_word = generate_random_word()
messages = [
{
"role": "user",
"content": f"write a one sentence poem about: {random_word}",
}
]
litellm.cache = Cache(
2024-03-08 00:56:59 +08:00
type="s3",
s3_bucket_name="litellm-proxy",
s3_region_name="us-west-2",
)
print("s3 Cache: test for caching, streaming + completion")
2024-01-03 17:43:39 +08:00
response_1_content = ""
response_2_content = ""
response_1_created = ""
response_2_created = ""
if sync_mode:
response1 = litellm.completion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
for chunk in response1:
print(chunk)
response_1_created = chunk.created
response_1_content += chunk.choices[0].delta.content or ""
print(response_1_content)
else:
response1 = await litellm.acompletion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response1:
print(chunk)
response_1_created = chunk.created
response_1_content += chunk.choices[0].delta.content or ""
print(response_1_content)
if sync_mode:
time.sleep(0.5)
else:
await asyncio.sleep(0.5)
2024-01-03 17:43:39 +08:00
print("\n\n Response 1 content: ", response_1_content, "\n\n")
if sync_mode:
response2 = litellm.completion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
for chunk in response2:
print(chunk)
response_2_content += chunk.choices[0].delta.content or ""
response_2_created = chunk.created
print(response_2_content)
else:
response2 = await litellm.acompletion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
messages=messages,
max_tokens=40,
temperature=1,
stream=True,
)
async for chunk in response2:
print(chunk)
response_2_content += chunk.choices[0].delta.content or ""
response_2_created = chunk.created
print(response_2_content)
2024-03-16 09:09:25 +08:00
2024-01-03 17:43:39 +08:00
print("\nresponse 1", response_1_content)
print("\nresponse 2", response_2_content)
2024-01-03 17:43:39 +08:00
assert (
response_1_content == response_2_content
), f"Response 1 != Response 2. Same params, Response 1{response_1_content} != Response 2{response_2_content}"
2024-01-05 19:10:56 +08:00
# prioritizing getting a new deploy out - will look at this in the next deploy
# print("response 1 created", response_1_created)
# print("response 2 created", response_2_created)
2024-01-05 19:10:56 +08:00
# assert response_1_created == response_2_created
2024-01-03 17:43:39 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
except Exception as e:
print(e)
raise e
# test_s3_cache_acompletion_stream_azure()
2024-01-03 17:43:39 +08:00
@pytest.mark.skip(reason="AWS Suspended Account")
2024-02-09 02:04:10 +08:00
@pytest.mark.asyncio
async def test_s3_cache_acompletion_azure():
import asyncio
import logging
import tracemalloc
tracemalloc.start()
logging.basicConfig(level=logging.DEBUG)
try:
litellm.set_verbose = True
random_word = generate_random_word()
messages = [
{
"role": "user",
"content": f"write a one sentence poem about: {random_word}",
}
]
litellm.cache = Cache(
type="s3",
s3_bucket_name="litellm-my-test-bucket-2",
s3_region_name="us-east-1",
2024-02-09 02:04:10 +08:00
)
print("s3 Cache: test for caching, streaming + completion")
response1 = await litellm.acompletion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
2024-02-09 02:04:10 +08:00
messages=messages,
max_tokens=40,
temperature=1,
)
print(response1)
time.sleep(2)
response2 = await litellm.acompletion(
2025-10-26 01:19:24 +08:00
model="azure/gpt-4.1-mini",
2024-02-09 02:04:10 +08:00
messages=messages,
max_tokens=40,
temperature=1,
)
print(response2)
assert response1.id == response2.id
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
except Exception as e:
print(e)
raise e
2023-12-12 00:43:00 +08:00
# test_redis_cache_acompletion_stream_bedrock()
2023-09-09 09:06:46 +08:00
# redis cache with custom keys
def custom_get_cache_key(*args, **kwargs):
2023-12-25 16:40:38 +08:00
# return key to use for your cache:
key = (
kwargs.get("model", "")
+ str(kwargs.get("messages", ""))
+ str(kwargs.get("temperature", ""))
+ str(kwargs.get("logit_bias", ""))
)
2023-09-09 09:06:46 +08:00
return key
2023-12-25 16:40:38 +08:00
2023-09-09 09:06:46 +08:00
def test_custom_redis_cache_with_key():
2023-09-09 11:15:14 +08:00
messages = [{"role": "user", "content": "write a one line story"}]
2023-12-25 16:40:38 +08:00
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
2023-09-09 09:06:46 +08:00
litellm.cache.get_cache_key = custom_get_cache_key
2023-09-12 01:36:17 +08:00
local_cache = {}
def set_cache(key, value):
local_cache[key] = value
2023-10-30 20:35:55 +08:00
2023-09-12 01:36:17 +08:00
def get_cache(key):
if key in local_cache:
return local_cache[key]
2023-10-30 20:35:55 +08:00
2023-09-12 01:36:17 +08:00
litellm.cache.cache.set_cache = set_cache
litellm.cache.cache.get_cache = get_cache
# patch this redis cache get and set call
2023-12-25 16:40:38 +08:00
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
temperature=1,
caching=True,
num_retries=3,
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
temperature=1,
caching=True,
num_retries=3,
)
response3 = completion(
model="gpt-3.5-turbo",
messages=messages,
temperature=1,
caching=False,
num_retries=3,
)
2023-10-30 20:35:55 +08:00
2023-09-09 09:06:46 +08:00
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
2023-12-25 16:40:38 +08:00
if (
response3["choices"][0]["message"]["content"]
== response2["choices"][0]["message"]["content"]
):
2023-10-30 20:35:55 +08:00
pytest.fail(f"Error occurred:")
2023-10-25 06:02:18 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-09-09 11:15:14 +08:00
2023-12-25 16:40:38 +08:00
2023-10-30 20:35:55 +08:00
# test_custom_redis_cache_with_key()
2023-08-29 13:10:15 +08:00
2023-12-25 16:40:38 +08:00
def test_cache_override():
# test if we can override the cache, when `caching=False` but litellm.cache = Cache() is set
2023-12-25 16:40:38 +08:00
# in this case it should not return cached responses
litellm.cache = Cache()
print("Testing cache override")
2023-12-25 16:40:38 +08:00
litellm.set_verbose = True
# test embedding
response1 = embedding(
2023-12-25 16:40:38 +08:00
model="text-embedding-ada-002", input=["hello who are you"], caching=False
)
start_time = time.time()
response2 = embedding(
2023-12-25 16:40:38 +08:00
model="text-embedding-ada-002", input=["hello who are you"], caching=False
)
end_time = time.time()
print(f"Embedding 2 response time: {end_time - start_time} seconds")
2023-12-25 16:40:38 +08:00
assert (
end_time - start_time > 0.05
) # ensure 2nd response comes in over 0.05s. This should not be cached.
2023-12-25 16:40:38 +08:00
# test_cache_override()
2023-12-05 12:50:06 +08:00
2024-03-26 09:56:36 +08:00
@pytest.mark.asyncio
async def test_cache_control_overrides():
# we use the cache controls to ensure there is no cache hit on this test
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
print("Testing cache override")
litellm.set_verbose = True
from litellm._uuid import uuid
2024-03-26 09:56:36 +08:00
unique_num = str(uuid.uuid4())
start_time = time.time()
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "hello who are you" + unique_num,
}
],
caching=True,
2024-03-26 09:56:36 +08:00
)
print(response1)
await asyncio.sleep(2)
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "hello who are you" + unique_num,
}
],
caching=True,
cache={"no-cache": True},
)
print(response2)
assert response1.id != response2.id
def test_sync_cache_control_overrides():
# we use the cache controls to ensure there is no cache hit on this test
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
print("Testing cache override")
litellm.set_verbose = True
from litellm._uuid import uuid
unique_num = str(uuid.uuid4())
start_time = time.time()
response1 = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "hello who are you" + unique_num,
}
],
caching=True,
)
print(response1)
time.sleep(2)
response2 = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "hello who are you" + unique_num,
}
],
caching=True,
2024-03-26 09:56:36 +08:00
cache={"no-cache": True},
)
print(response2)
assert response1.id != response2.id
2023-12-05 12:50:06 +08:00
def test_custom_redis_cache_params():
# test if we can init redis with **kwargs
try:
litellm.cache = Cache(
type="redis",
2023-12-25 16:40:38 +08:00
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
db=0,
2023-12-05 12:50:06 +08:00
)
2023-12-25 16:40:38 +08:00
print(litellm.cache.cache.redis_client)
2023-12-05 12:50:06 +08:00
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
2023-12-05 12:50:06 +08:00
except Exception as e:
2024-04-20 08:02:15 +08:00
pytest.fail(f"Error occurred: {str(e)}")
2023-12-05 12:50:06 +08:00
2023-12-10 08:42:54 +08:00
def test_get_cache_key():
from litellm.caching.caching import Cache
2023-12-25 16:40:38 +08:00
2023-12-10 08:42:54 +08:00
try:
print("Testing get_cache_key")
2023-12-10 08:42:54 +08:00
cache_instance = Cache()
2023-12-25 16:40:38 +08:00
cache_key = cache_instance.get_cache_key(
**{
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "write a one sentence poem about: 7510"}
],
"max_tokens": 40,
"temperature": 0.2,
"stream": True,
"litellm_call_id": "ffe75e7e-8a07-431f-9a74-71a5b9f35f0b",
"litellm_logging_obj": {},
}
2023-12-10 08:42:54 +08:00
)
2023-12-25 16:40:38 +08:00
cache_key_2 = cache_instance.get_cache_key(
**{
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "write a one sentence poem about: 7510"}
],
"max_tokens": 40,
"temperature": 0.2,
"stream": True,
"litellm_call_id": "ffe75e7e-8a07-431f-9a74-71a5b9f35f0b",
"litellm_logging_obj": {},
}
)
cache_key_str = "model: gpt-3.5-turbomessages: [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}]max_tokens: 40temperature: 0.2stream: True"
hash_object = hashlib.sha256(cache_key_str.encode())
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
assert cache_key == hash_hex
2023-12-25 16:40:38 +08:00
assert (
cache_key_2 == hash_hex
2023-12-25 16:40:38 +08:00
), f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs"
embedding_cache_key = cache_instance.get_cache_key(
2023-12-25 16:40:38 +08:00
**{
2025-09-28 03:41:35 +08:00
"model": "azure/text-embedding-ada-002",
2023-12-25 16:40:38 +08:00
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_key": "",
"api_version": "2023-07-01-preview",
"timeout": None,
"max_retries": 0,
"input": ["hi who is ishaan"],
"caching": True,
"client": "<openai.lib.azure.AsyncAzureOpenAI object at 0x12b6a1060>",
}
)
print(embedding_cache_key)
embedding_cache_key_str = (
2025-09-28 03:41:35 +08:00
"model: azure/text-embedding-ada-002input: ['hi who is ishaan']"
)
hash_object = hashlib.sha256(embedding_cache_key_str.encode())
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
2023-12-25 16:40:38 +08:00
assert (
embedding_cache_key == hash_hex
2025-09-28 03:41:35 +08:00
), f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs"
# Proxy - embedding cache, test if embedding key, gets model_group and not model
embedding_cache_key_2 = cache_instance.get_cache_key(
2023-12-25 16:40:38 +08:00
**{
2025-09-28 03:41:35 +08:00
"model": "azure/text-embedding-ada-002",
2023-12-25 16:40:38 +08:00
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_key": "",
"api_version": "2023-07-01-preview",
"timeout": None,
"max_retries": 0,
"input": ["hi who is ishaan"],
"caching": True,
"client": "<openai.lib.azure.AsyncAzureOpenAI object at 0x12b6a1060>",
"proxy_server_request": {
"url": "http://0.0.0.0:8000/embeddings",
"method": "POST",
"headers": {
"host": "0.0.0.0:8000",
"user-agent": "curl/7.88.1",
"accept": "*/*",
"content-type": "application/json",
"content-length": "80",
},
"body": {
"model": "azure-embedding-model",
"input": ["hi who is ishaan"],
},
},
"user": None,
"metadata": {
"user_api_key": None,
"headers": {
"host": "0.0.0.0:8000",
"user-agent": "curl/7.88.1",
"accept": "*/*",
"content-type": "application/json",
"content-length": "80",
},
"model_group": "EMBEDDING_MODEL_GROUP",
2025-09-28 03:41:35 +08:00
"deployment": "azure/text-embedding-ada-002-ModelID-azure/text-embedding-ada-002https://openai-gpt-4-test-v-1.openai.azure.com/2023-07-01-preview",
2023-12-25 16:40:38 +08:00
},
"model_info": {
"mode": "embedding",
"base_model": "text-embedding-ada-002",
"id": "20b2b515-f151-4dd5-a74f-2231e2f54e29",
},
"litellm_call_id": "2642e009-b3cd-443d-b5dd-bb7d56123b0e",
"litellm_logging_obj": "<litellm.utils.Logging object at 0x12f1bddb0>",
}
)
print(embedding_cache_key_2)
embedding_cache_key_str_2 = (
"model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']"
2023-12-25 16:40:38 +08:00
)
hash_object = hashlib.sha256(embedding_cache_key_str_2.encode())
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
assert embedding_cache_key_2 == hash_hex
print("passed!")
2023-12-10 08:42:54 +08:00
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred:", e)
2023-12-25 16:40:38 +08:00
# test_get_cache_key()
2023-12-10 08:42:54 +08:00
2023-12-30 22:03:47 +08:00
def test_cache_context_managers():
litellm.set_verbose = True
litellm.cache = Cache(type="redis")
# cache is on, disable it
litellm.disable_cache()
assert litellm.cache == None
assert "cache" not in litellm.success_callback
assert "cache" not in litellm._async_success_callback
# disable a cache that is off
litellm.disable_cache()
assert litellm.cache == None
assert "cache" not in litellm.success_callback
assert "cache" not in litellm._async_success_callback
litellm.enable_cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
)
assert litellm.cache != None
assert litellm.cache.type == "redis"
print("VARS of litellm.cache", vars(litellm.cache))
2024-02-06 04:28:21 +08:00
def test_redis_semantic_cache_completion():
2024-02-06 09:58:32 +08:00
litellm.set_verbose = True
2024-02-07 00:14:54 +08:00
import logging
logging.basicConfig(level=logging.DEBUG)
2024-02-06 04:28:21 +08:00
2024-02-06 10:22:50 +08:00
print("testing semantic caching")
2024-02-06 04:28:21 +08:00
litellm.cache = Cache(
type="redis-semantic",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
2024-02-06 10:22:50 +08:00
similarity_threshold=0.8,
2024-02-07 02:55:15 +08:00
redis_semantic_cache_embedding_model="text-embedding-ada-002",
2024-02-06 04:28:21 +08:00
)
2024-02-06 10:22:50 +08:00
response1 = completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
2025-03-19 10:05:26 +08:00
"content": "write a one sentence poem about summer",
2024-02-06 10:22:50 +08:00
}
],
max_tokens=20,
)
print(f"response1: {response1}")
2024-02-07 02:39:44 +08:00
response2 = completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
2025-03-19 10:05:26 +08:00
"content": "write a one sentence poem about summertime",
2024-02-07 02:39:44 +08:00
}
],
max_tokens=20,
)
2025-03-19 10:05:26 +08:00
print(f"response2: {response2}")
2024-02-07 02:39:44 +08:00
assert response1.id == response2.id
2024-02-06 04:28:21 +08:00
# test_redis_cache_completion()
2024-02-07 00:14:54 +08:00
2025-03-25 14:15:04 +08:00
@pytest.mark.flaky(reruns=3)
2024-02-07 00:14:54 +08:00
@pytest.mark.asyncio
async def test_redis_semantic_cache_acompletion():
litellm.set_verbose = True
import logging
logging.basicConfig(level=logging.DEBUG)
print("testing semantic caching")
litellm.cache = Cache(
type="redis-semantic",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
2025-03-19 10:05:26 +08:00
similarity_threshold=0.7,
2024-02-07 00:14:54 +08:00
)
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
2025-03-19 10:05:26 +08:00
"content": "write a one sentence poem about summer",
2024-02-07 00:14:54 +08:00
}
],
2024-02-07 02:39:44 +08:00
max_tokens=5,
2024-02-07 00:14:54 +08:00
)
print(f"response1: {response1}")
2025-03-25 14:15:04 +08:00
await asyncio.sleep(2)
2024-02-07 02:39:44 +08:00
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
2025-03-19 10:05:26 +08:00
"content": "write a one sentence poem about summertime",
2024-02-07 02:39:44 +08:00
}
],
max_tokens=5,
)
print(f"response2: {response2}")
assert response1.id == response2.id
def test_caching_redis_simple(caplog, capsys):
"""
Relevant issue - https://github.com/BerriAI/litellm/issues/4511
"""
litellm.set_verbose = True ## REQUIRED FOR TEST.
litellm.cache = Cache(
type="redis", url=os.getenv("REDIS_SSL_URL")
) # passing `supported_call_types = ["completion"]` has no effect
s = time.time()
uuid_str = str(uuid.uuid4())
x = completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": f"Hello, how are you? Wink {uuid_str}"}],
stream=True,
)
for m in x:
print(m)
print(time.time() - s)
s2 = time.time()
x = completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": f"Hello, how are you? Wink {uuid_str}"}],
stream=True,
)
for m in x:
print(m)
print(time.time() - s2)
redis_async_caching_error = False
redis_service_logging_error = False
captured = capsys.readouterr()
captured_logs = [rec.message for rec in caplog.records]
print(f"captured_logs: {captured_logs}")
for item in captured_logs:
if (
"Error connecting to Async Redis client" in item
or "Set ASYNC Redis Cache" in item
):
redis_async_caching_error = True
if "ServiceLogging.async_service_success_hook" in item:
redis_service_logging_error = True
assert redis_async_caching_error is False
assert redis_service_logging_error is False
assert "async success_callback: reaches cache for logging" not in captured.out
2024-08-22 03:07:57 +08:00
2024-08-25 00:06:59 +08:00
@pytest.mark.asyncio()
async def test_cache_default_off_acompletion():
litellm.set_verbose = True
import logging
from litellm._logging import verbose_logger
verbose_logger.setLevel(logging.DEBUG)
from litellm.caching.caching import CacheMode
2024-08-25 00:06:59 +08:00
random_number = random.randint(
1, 100000
) # add a random number to ensure it's always adding /reading from cache
litellm.cache = Cache(
type="local",
mode=CacheMode.default_off,
)
### No Cache hits when it's default off
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": f"write a one sentence poem about: {random_number}",
}
],
mock_response="hello",
max_tokens=20,
)
print(f"Response1: {response1}")
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": f"write a one sentence poem about: {random_number}",
}
],
max_tokens=20,
)
print(f"Response2: {response2}")
assert response1.id != response2.id
## Cache hits when it's default off and then opt in
response3 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": f"write a one sentence poem about: {random_number}",
}
],
mock_response="hello",
cache={"use-cache": True},
metadata={"key": "value"},
max_tokens=20,
)
print(f"Response3: {response3}")
2024-08-25 00:37:41 +08:00
await asyncio.sleep(2)
2024-08-25 00:06:59 +08:00
response4 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": f"write a one sentence poem about: {random_number}",
}
],
cache={"use-cache": True},
metadata={"key": "value"},
max_tokens=20,
)
print(f"Response4: {response4}")
assert response3.id == response4.id
2024-09-04 22:57:23 +08:00
LiteLLM Minor Fixes & Improvements (09/18/2024) (#5772) * fix(proxy_server.py): fix azure key vault logic to not require client id/secret * feat(cost_calculator.py): support fireworks ai cost tracking * build(docker-compose.yml): add lines for mounting config.yaml to docker compose Closes https://github.com/BerriAI/litellm/issues/5739 * fix(input.md): update docs to clarify litellm supports content as a list of dictionaries Fixes https://github.com/BerriAI/litellm/issues/5755 * fix(input.md): update input.md to include all message values * fix(image_handling.py): follow image url redirects Fixes https://github.com/BerriAI/litellm/issues/5763 * fix(router.py): Fix model key/base leak in error message Fixes https://github.com/BerriAI/litellm/issues/5762 * fix(http_handler.py): fix linting error * fix(azure.py): fix logging to show azure_ad_token being used Fixes https://github.com/BerriAI/litellm/issues/5767 * fix(_redis.py): add redis sentinel support Closes https://github.com/BerriAI/litellm/issues/4381 * feat(_redis.py): add redis sentinel support Closes https://github.com/BerriAI/litellm/issues/4381 * test(test_completion_cost.py): fix test * Databricks Integration: Integrate Databricks SDK as optional mechanism for fetching API base and token, if unspecified (#5746) * LiteLLM Minor Fixes & Improvements (09/16/2024) (#5723) * coverage (#5713) Signed-off-by: dbczumar <corey.zumar@databricks.com> * Move (#5714) Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix(litellm_logging.py): fix logging client re-init (#5710) Fixes https://github.com/BerriAI/litellm/issues/5695 * fix(presidio.py): Fix logging_hook response and add support for additional presidio variables in guardrails config Fixes https://github.com/BerriAI/litellm/issues/5682 * feat(o1_handler.py): fake streaming for openai o1 models Fixes https://github.com/BerriAI/litellm/issues/5694 * docs: deprecated traceloop integration in favor of native otel (#5249) * fix: fix linting errors * fix: fix linting errors * fix(main.py): fix o1 import --------- Signed-off-by: dbczumar <corey.zumar@databricks.com> Co-authored-by: Corey Zumar <39497902+dbczumar@users.noreply.github.com> Co-authored-by: Nir Gazit <nirga@users.noreply.github.com> * feat(spend_management_endpoints.py): expose `/global/spend/refresh` endpoint for updating material view (#5730) * feat(spend_management_endpoints.py): expose `/global/spend/refresh` endpoint for updating material view Supports having `MonthlyGlobalSpend` view be a material view, and exposes an endpoint to refresh it * fix(custom_logger.py): reset calltype * fix: fix linting errors * fix: fix linting error * fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix: fix import * Fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * DB test Signed-off-by: dbczumar <corey.zumar@databricks.com> * Coverage Signed-off-by: dbczumar <corey.zumar@databricks.com> * progress Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * fix test name Signed-off-by: dbczumar <corey.zumar@databricks.com> --------- Signed-off-by: dbczumar <corey.zumar@databricks.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com> Co-authored-by: Nir Gazit <nirga@users.noreply.github.com> * test: fix test * test(test_databricks.py): fix test * fix(databricks/chat.py): handle custom endpoint (e.g. sagemaker) * Apply code scanning fix for clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * fix(__init__.py): fix known fireworks ai models --------- Signed-off-by: dbczumar <corey.zumar@databricks.com> Co-authored-by: Corey Zumar <39497902+dbczumar@users.noreply.github.com> Co-authored-by: Nir Gazit <nirga@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
2024-09-20 04:25:29 +08:00
@pytest.mark.skip(reason="local test. Requires sentinel setup.")
@pytest.mark.asyncio
async def test_redis_sentinel_caching():
"""
Init redis client
- write to client
- read from client
"""
litellm.set_verbose = False
random_number = random.randint(
1, 100000
) # add a random number to ensure it's always adding / reading from cache
messages = [
{"role": "user", "content": f"write a one sentence poem about: {random_number}"}
]
litellm.cache = Cache(
type="redis",
# host=os.environ["REDIS_HOST"],
# port=os.environ["REDIS_PORT"],
# password=os.environ["REDIS_PASSWORD"],
service_name="mymaster",
sentinel_nodes=[("localhost", 26379)],
)
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
)
cache_key = litellm.cache.get_cache_key(
model="gpt-3.5-turbo",
messages=messages,
)
print(f"cache_key: {cache_key}")
litellm.cache.add_cache(result=response1, cache_key=cache_key)
print(f"cache key pre async get: {cache_key}")
stored_val = litellm.cache.get_cache(
model="gpt-3.5-turbo",
messages=messages,
)
print(f"stored_val: {stored_val}")
assert stored_val["id"] == response1.id
stored_val_2 = await litellm.cache.async_get_cache(
model="gpt-3.5-turbo",
messages=messages,
)
print(f"stored_val: {stored_val}")
assert stored_val_2["id"] == response1.id
@pytest.mark.asyncio
2025-07-20 06:58:25 +08:00
@pytest.mark.flaky(retries=3, delay=2)
async def test_redis_proxy_batch_redis_get_cache():
"""
Tests batch_redis_get.py
- make 1st call -> expect miss
- make 2nd call -> expect hit
"""
from litellm.caching.caching import Cache, DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.batch_redis_get import _PROXY_BatchRedisRequests
litellm.cache = Cache(
type="redis",
host=os.getenv("REDIS_HOST"),
port=os.getenv("REDIS_PORT"),
password=os.getenv("REDIS_PASSWORD"),
namespace="test_namespace",
)
batch_redis_get_obj = (
_PROXY_BatchRedisRequests()
) # overrides the .async_get_cache method
user_api_key_cache = DualCache()
from litellm._uuid import uuid
batch_redis_get_obj.in_memory_cache = user_api_key_cache.in_memory_cache
messages = [{"role": "user", "content": "hi {}".format(uuid.uuid4())}]
# 1st call -> expect miss
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
mock_response="hello",
)
assert response is not None
assert "cache_key" not in response._hidden_params
print(response._hidden_params)
await asyncio.sleep(1)
# 2nd call -> expect hit
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
mock_response="hello",
)
print(response._hidden_params)
assert "cache_key" in response._hidden_params
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_logging_turn_off_message_logging_streaming(sync_mode):
litellm.turn_off_message_logging = True
mock_obj = Cache(type="local")
litellm.cache = mock_obj
with patch.object(mock_obj, "add_cache") as mock_client, patch.object(
mock_obj, "async_add_cache"
) as mock_async_client:
print(f"mock_obj.add_cache: {mock_obj.add_cache}")
if sync_mode is True:
resp = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
mock_response="hello",
stream=True,
)
for chunk in resp:
continue
time.sleep(1)
mock_client.assert_called_once()
print(f"mock_client.call_args: {mock_client.call_args}")
assert mock_client.call_args.args[0].choices[0].message.content == "hello"
else:
resp = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
mock_response="hello",
stream=True,
)
async for chunk in resp:
continue
await asyncio.sleep(1)
mock_async_client.assert_called_once()
print(f"mock_async_client.call_args: {mock_async_client.call_args.args[0]}")
print(
f"mock_async_client.call_args: {json.loads(mock_async_client.call_args.args[0])}"
)
json_mock = json.loads(mock_async_client.call_args.args[0])
try:
assert json_mock["choices"][0]["message"]["content"] == "hello"
except Exception as e:
print(
f"mock_async_client.call_args.args[0]: {mock_async_client.call_args.args[0]}"
)
print(
f"mock_async_client.call_args.args[0]['choices']: {mock_async_client.call_args.args[0]['choices']}"
)
raise e
Litellm Minor Fixes & Improvements (10/12/2024) (#6179) * build(model_prices_and_context_window.json): add bedrock llama3.2 pricing * build(model_prices_and_context_window.json): add bedrock cross region inference pricing * Revert "(perf) move s3 logging to Batch logging + async [94% faster perf under 100 RPS on 1 litellm instance] (#6165)" This reverts commit 2a5624af471284f174e084142504d950ede2567d. * add azure/gpt-4o-2024-05-13 (#6174) * LiteLLM Minor Fixes & Improvements (10/10/2024) (#6158) * refactor(vertex_ai_partner_models/anthropic): refactor anthropic to use partner model logic * fix(vertex_ai/): support passing custom api base to partner models Fixes https://github.com/BerriAI/litellm/issues/4317 * fix(proxy_server.py): Fix prometheus premium user check logic * docs(prometheus.md): update quick start docs * fix(custom_llm.py): support passing dynamic api key + api base * fix(realtime_api/main.py): Add request/response logging for realtime api endpoints Closes https://github.com/BerriAI/litellm/issues/6081 * feat(openai/realtime): add openai realtime api logging Closes https://github.com/BerriAI/litellm/issues/6081 * fix(realtime_streaming.py): fix linting errors * fix(realtime_streaming.py): fix linting errors * fix: fix linting errors * fix pattern match router * Add literalai in the sidebar observability category (#6163) * fix: add literalai in the sidebar * fix: typo * update (#6160) * Feat: Add Langtrace integration (#5341) * Feat: Add Langtrace integration * add langtrace service name * fix timestamps for traces * add tests * Discard Callback + use existing otel logger * cleanup * remove print statments * remove callback * add docs * docs * add logging docs * format logging * remove emoji and add litellm proxy example * format logging * format `logging.md` * add langtrace docs to logging.md * sync conflict * docs fix * (perf) move s3 logging to Batch logging + async [94% faster perf under 100 RPS on 1 litellm instance] (#6165) * fix move s3 to use customLogger * add basic s3 logging test * add s3 to custom logger compatible * use batch logger for s3 * s3 set flush interval and batch size * fix s3 logging * add notes on s3 logging * fix s3 logging * add basic s3 logging test * fix s3 type errors * add test for sync logging on s3 * fix: fix to debug log --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Willy Douhard <willy.douhard@gmail.com> Co-authored-by: yujonglee <yujonglee.dev@gmail.com> Co-authored-by: Ali Waleed <ali@scale3labs.com> * docs(custom_llm_server.md): update doc on passing custom params * fix(pass_through_endpoints.py): don't require headers Fixes https://github.com/BerriAI/litellm/issues/6128 * feat(utils.py): add support for caching rerank endpoints Closes https://github.com/BerriAI/litellm/issues/6144 * feat(litellm_logging.py'): add response headers for failed requests Closes https://github.com/BerriAI/litellm/issues/6159 --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Willy Douhard <willy.douhard@gmail.com> Co-authored-by: yujonglee <yujonglee.dev@gmail.com> Co-authored-by: Ali Waleed <ali@scale3labs.com>
2024-10-13 02:48:34 +08:00
def test_basic_caching_import():
from litellm.caching import Cache
assert Cache is not None
print("Cache imported successfully")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio()
async def test_caching_kwargs_input(sync_mode):
from litellm import acompletion
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.types.utils import (
Choices,
EmbeddingResponse,
Message,
ModelResponse,
Usage,
CompletionTokensDetailsWrapper,
PromptTokensDetailsWrapper,
)
from datetime import datetime
llm_caching_handler = LLMCachingHandler(
original_function=acompletion, request_kwargs={}, start_time=datetime.now()
)
input = {
"result": ModelResponse(
id="chatcmpl-AJ119H5XsDnYiZPp5axJ5d7niwqeR",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Hello! I'm just a computer program, so I don't have feelings, but I'm here to assist you. How can I help you today?",
role="assistant",
tool_calls=None,
function_call=None,
),
)
],
created=1729095507,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint=None,
usage=Usage(
completion_tokens=31,
prompt_tokens=16,
total_tokens=47,
completion_tokens_details=CompletionTokensDetailsWrapper(
audio_tokens=None, reasoning_tokens=0
),
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None, cached_tokens=0
),
),
service_tier=None,
),
"kwargs": {
"messages": [{"role": "user", "content": "42HHey, how's it going?"}],
"caching": True,
"litellm_call_id": "fae2aa4f-9f75-4f11-8c9c-63ab8d9fae26",
"preset_cache_key": "2f69f5640d5e0f25315d0e132f1278bb643554d14565d2c61d61564b10ade90f",
},
"args": ("gpt-3.5-turbo",),
}
if sync_mode is True:
llm_caching_handler.sync_set_cache(**input)
else:
input["original_function"] = acompletion
await llm_caching_handler.async_set_cache(**input)
@pytest.mark.skip(reason="audio caching not supported yet")
@pytest.mark.parametrize("stream", [False]) # True,
@pytest.mark.asyncio()
async def test_audio_caching(stream):
litellm.cache = Cache(type="local")
## CALL 1 - no cache hit
completion = await litellm.acompletion(
model="gpt-4o-audio-preview",
modalities=["text", "audio"],
audio={"voice": "alloy", "format": "pcm16"},
messages=[{"role": "user", "content": "response in 1 word - yes or no"}],
stream=stream,
)
assert "cache_hit" not in completion._hidden_params
## CALL 2 - cache hit
completion = await litellm.acompletion(
model="gpt-4o-audio-preview",
modalities=["text", "audio"],
audio={"voice": "alloy", "format": "pcm16"},
messages=[{"role": "user", "content": "response in 1 word - yes or no"}],
stream=stream,
)
assert "cache_hit" in completion._hidden_params
def test_redis_caching_default_ttl():
"""
Ensure that the default redis cache TTL is 60s
"""
from litellm.caching.redis_cache import RedisCache
litellm.default_redis_ttl = 120
cache_obj = RedisCache()
assert cache_obj.default_ttl == 120
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_redis_caching_llm_caching_ttl(sync_mode):
"""
Ensure default redis cache ttl is used for a sample redis cache object
"""
from litellm.caching.redis_cache import RedisCache
litellm.default_redis_ttl = 120
cache_obj = RedisCache()
assert cache_obj.default_ttl == 120
if sync_mode is False:
# Create an AsyncMock for the Redis client
mock_redis_instance = AsyncMock()
# Make sure the mock can be used as an async context manager
mock_redis_instance.__aenter__.return_value = mock_redis_instance
mock_redis_instance.__aexit__.return_value = None
## Set cache
if sync_mode is True:
with patch.object(cache_obj.redis_client, "set") as mock_set:
cache_obj.set_cache(key="test", value="test")
mock_set.assert_called_once_with(name="test", value="test", ex=120)
else:
# Patch self.init_async_client to return our mock Redis client
with patch.object(
cache_obj, "init_async_client", return_value=mock_redis_instance
):
# Call async_set_cache
await cache_obj.async_set_cache(key="test", value="test_value")
# Verify that the set method was called on the mock Redis instance
mock_redis_instance.set.assert_called_once_with(
2025-04-03 12:54:35 +08:00
name="test", value='"test_value"', ex=120, nx=False
)
## Increment cache
if sync_mode is True:
with patch.object(cache_obj.redis_client, "ttl") as mock_incr:
cache_obj.increment_cache(key="test", value=1)
mock_incr.assert_called_once_with("test")
else:
# Patch self.init_async_client to return our mock Redis client
with patch.object(
cache_obj, "init_async_client", return_value=mock_redis_instance
):
# Call async_set_cache
await cache_obj.async_increment(key="test", value="test_value")
# Verify that the set method was called on the mock Redis instance
mock_redis_instance.ttl.assert_called_once_with("test")
@pytest.mark.asyncio()
async def test_redis_caching_ttl_pipeline():
"""
Ensure that a default ttl is set for all redis functions
"""
from litellm.caching.redis_cache import RedisCache
litellm.default_redis_ttl = 120
expected_timedelta = timedelta(seconds=120)
cache_obj = RedisCache()
## TEST 1 - async_set_cache_pipeline
# Patch self.init_async_client to return our mock Redis client
# Call async_set_cache
mock_pipe_instance = AsyncMock()
with patch.object(mock_pipe_instance, "set", return_value=None) as mock_set:
await cache_obj._pipeline_helper(
pipe=mock_pipe_instance,
cache_list=[("test_key1", "test_value1"), ("test_key2", "test_value2")],
ttl=None,
)
# Verify that the set method was called on the mock Redis instance
mock_set.assert_has_calls(
[
call.set(
name="test_key1", value='"test_value1"', ex=expected_timedelta
),
call.set(
name="test_key2", value='"test_value2"', ex=expected_timedelta
),
]
)
@pytest.mark.asyncio()
async def test_redis_caching_ttl_sadd():
"""
Ensure that a default ttl is set for all redis functions
"""
from litellm.caching.redis_cache import RedisCache
litellm.default_redis_ttl = 120
expected_timedelta = timedelta(seconds=120)
cache_obj = RedisCache()
redis_client = AsyncMock()
with patch.object(redis_client, "expire", return_value=None) as mock_expire:
await cache_obj._set_cache_sadd_helper(
redis_client=redis_client, key="test_key", value=["test_value"], ttl=None
)
print(f"expected_timedelta: {expected_timedelta}")
assert mock_expire.call_args.args[1] == expected_timedelta
(perf) Litellm redis router fix - ~100ms improvement (#6483) * docs(exception_mapping.md): add missing exception types Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183 * fix(main.py): register custom model pricing with specific key Ensure custom model pricing is registered to the specific model+provider key combination * test: make testing more robust for custom pricing * fix(redis_cache.py): instrument otel logging for sync redis calls ensures complete coverage for all redis cache calls * refactor: pass parent_otel_span for redis caching calls in router allows for more observability into what calls are causing latency issues * test: update tests with new params * refactor: ensure e2e otel tracing for router * refactor(router.py): add more otel tracing acrosss router catch all latency issues for router requests * fix: fix linting error * fix(router.py): fix linting error * fix: fix test * test: fix tests * fix(dual_cache.py): pass ttl to redis cache * fix: fix param * perf(cooldown_cache.py): improve cooldown cache, to store cache results in memory for 5s, prevents redis call from being made on each request reduces 100ms latency per call with caching enabled on router * fix: fix test * fix(cooldown_cache.py): handle if a result is None * fix(cooldown_cache.py): add debug statements * refactor(dual_cache.py): move to using an in-memory check for batch get cache, to prevent redis from being hit for every call * fix(cooldown_cache.py): fix linting erropr
2024-10-30 04:58:29 +08:00
@pytest.mark.asyncio()
async def test_dual_cache_caching_batch_get_cache():
"""
- check redis cache called for initial batch get cache
- check redis cache not called for consecutive batch get cache with same keys
"""
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache
dc = DualCache(redis_cache=MagicMock(spec=RedisCache))
with patch.object(
dc.redis_cache,
"async_batch_get_cache",
new=AsyncMock(
return_value={"test_key1": "test_value1", "test_key2": "test_value2"}
),
) as mock_async_get_cache:
await dc.async_batch_get_cache(keys=["test_key1", "test_key2"])
assert mock_async_get_cache.call_count == 1
await dc.async_batch_get_cache(keys=["test_key1", "test_key2"])
assert mock_async_get_cache.call_count == 1
@pytest.mark.asyncio
async def test_redis_increment_pipeline():
"""Test Redis increment pipeline functionality"""
try:
from litellm.caching.redis_cache import RedisCache
litellm.set_verbose = True
litellm._turn_on_debug()
redis_cache = RedisCache(
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
# Create test increment operations
increment_list = [
{"key": "test_key1", "increment_value": 1.5, "ttl": 60},
{"key": "test_key1", "increment_value": 1.1, "ttl": 58},
{"key": "test_key1", "increment_value": 0.4, "ttl": 55},
{"key": "test_key2", "increment_value": 2.5, "ttl": 60},
]
# Test pipeline increment
results = await redis_cache.async_increment_pipeline(increment_list)
# Verify results
assert len(results) == 4
# Verify the values were actually set in Redis
value1 = await redis_cache.async_get_cache("test_key1")
print("result in cache for key=test_key1", value1)
value2 = await redis_cache.async_get_cache("test_key2")
print("result in cache for key=test_key2", value2)
assert float(value1) == 3.0
assert float(value2) == 2.5
# Clean up
await redis_cache.async_delete_cache("test_key1")
await redis_cache.async_delete_cache("test_key2")
except Exception as e:
print(f"Error occurred: {str(e)}")
raise e
@pytest.mark.asyncio
async def test_redis_get_ttl():
"""
Test Redis get TTL functionality
Redis returns -2 if the key does not exist and -1 if the key exists but has no associated expire.
test that litellm redis caching wrapper handles -1 and -2 values and returns them as None
"""
try:
from litellm.caching.redis_cache import RedisCache
redis_cache = RedisCache(
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
# Test case 1: Key does not exist
result = await redis_cache.async_get_ttl("nonexistent_key")
print("ttl for nonexistent key: ", result)
assert result is None, f"Expected None for nonexistent key, got {result}"
# Test case 2: Key exists with TTL
test_key = "test_key_ttl"
test_value = "test_value"
ttl = 10 # 10 seconds TTL
# Set a key with TTL
_redis_client = await redis_cache.init_async_client()
async with _redis_client as redis_client:
await redis_client.set(test_key, test_value, ex=ttl)
# Get TTL and verify it's close to what we set
result = await redis_cache.async_get_ttl(test_key)
print("ttl for test_key: ", result)
assert (
result is not None and 0 <= result <= ttl
), f"Expected TTL between 0 and {ttl}, got {result}"
# Clean up
await redis_client.delete(test_key)
except Exception as e:
print(f"Error occurred: {str(e)}")
raise e
def test_redis_caching_multiple_namespaces():
"""
Test that redis caching works with multiple namespaces
If client side request specifies a namespace, it should be used for caching
The same request with different namespaces should not be cached under the same key
"""
from litellm._uuid import uuid
from unittest.mock import patch, MagicMock
import litellm
from litellm.caching import Cache
from litellm import completion
# Use a fixed uuid to ensure consistent cache keys
test_uuid = "12345678-1234-1234-1234-123456789abc"
messages = [{"role": "user", "content": f"what is litellm? {test_uuid}"}]
# Mock the Redis client creation from the _redis module
with patch("litellm._redis.get_redis_client") as mock_get_redis_client, patch(
"litellm._redis.get_redis_connection_pool"
) as mock_get_redis_connection_pool:
# Create a mock Redis client that simulates real Redis behavior
mock_redis_client = MagicMock()
mock_get_redis_client.return_value = mock_redis_client
# Mock the connection pool
mock_connection_pool = MagicMock()
mock_get_redis_connection_pool.return_value = mock_connection_pool
# Dictionary to simulate Redis storage with namespace support
redis_storage = {}
def mock_redis_get(key):
print(f"Redis GET: {key}")
value = redis_storage.get(key, None)
# Convert to bytes to match real Redis behavior
if value is not None:
import json
return json.dumps(value).encode("utf-8")
return None
def mock_redis_set(name, value, ex=None, **kwargs):
print(f"Redis SET: {name} = {value}")
redis_storage[name] = value
return True
def mock_redis_ping():
return True
def mock_redis_info():
return {"redis_version": "7.0.0"}
mock_redis_client.get = mock_redis_get
mock_redis_client.set = mock_redis_set
mock_redis_client.ping = mock_redis_ping
mock_redis_client.info = mock_redis_info
# Initialize the cache
litellm.cache = Cache(type="redis")
namespace_1 = "org-id1"
namespace_2 = "org-id2"
# Use mock_response to ensure deterministic responses without external API calls
response_1 = completion(
model="gpt-3.5-turbo",
messages=messages,
cache={"namespace": namespace_1},
mock_response="Response for namespace 1",
)
response_2 = completion(
model="gpt-3.5-turbo",
messages=messages,
cache={"namespace": namespace_2},
mock_response="Response for namespace 2",
)
response_3 = completion(
model="gpt-3.5-turbo",
messages=messages,
cache={"namespace": namespace_1},
mock_response="This should be cached",
)
response_4 = completion(
model="gpt-3.5-turbo",
messages=messages,
mock_response="Response without namespace",
)
print(
f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}"
)
print(
f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}"
)
print(
f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}"
)
print(
f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}"
)
print(f"Redis storage keys: {list(redis_storage.keys())}")
# Verify that different namespaces created different cache keys
cache_keys = list(redis_storage.keys())
namespace_1_keys = [k for k in cache_keys if k.startswith(f"{namespace_1}:")]
namespace_2_keys = [k for k in cache_keys if k.startswith(f"{namespace_2}:")]
no_namespace_keys = [
k
for k in cache_keys
if not k.startswith(f"{namespace_1}:")
and not k.startswith(f"{namespace_2}:")
]
print(f"Namespace 1 keys: {namespace_1_keys}")
print(f"Namespace 2 keys: {namespace_2_keys}")
print(f"No namespace keys: {no_namespace_keys}")
# Should have at least one key for each namespace
assert len(namespace_1_keys) > 0, "Should have cache keys for namespace 1"
assert len(namespace_2_keys) > 0, "Should have cache keys for namespace 2"
assert len(no_namespace_keys) > 0, "Should have cache keys for no namespace"
# The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1
assert isinstance(
response_3, str
), "Response 3 should be a cache hit (string) for same namespace"
# response 1 & 2 should be ModelResponse objects (cache misses)
assert hasattr(response_1, "id"), "Response 1 should be a ModelResponse object"
assert hasattr(response_2, "id"), "Response 2 should be a ModelResponse object"
assert hasattr(response_4, "id"), "Response 4 should be a ModelResponse object"
# response 1 & 2 should have different IDs (different namespaces)
assert (
response_1.id != response_2.id
), f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}"
# response 1 & 4 should have different IDs (different namespaces)
assert (
response_1.id != response_4.id
), f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}"
@pytest.mark.flaky(retries=3, delay=1)
def test_caching_with_reasoning_content():
"""
Test that reasoning content is cached
"""
from litellm._uuid import uuid
try:
messages = [{"role": "user", "content": f"what is litellm? {uuid.uuid4()}"}]
litellm.cache = Cache()
response_1 = completion(
model="anthropic/claude-3-7-sonnet-latest",
messages=messages,
thinking={"type": "enabled", "budget_tokens": 1024},
)
response_2 = completion(
model="anthropic/claude-3-7-sonnet-latest",
messages=messages,
thinking={"type": "enabled", "budget_tokens": 1024},
)
print(f"response 2: {response_2.model_dump_json(indent=4)}")
assert response_2._hidden_params["cache_hit"] == True
assert response_2.choices[0].message.reasoning_content is not None
except litellm.InternalServerError as e:
pytest.skip(f"Anthropic API returned InternalServerError - {str(e)}")
def test_caching_reasoning_args_miss(): # test in memory cache
try:
# litellm._turn_on_debug()
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
reasoning_effort="low",
mock_response="My response",
)
response2 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
mock_response="My response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response1.id != response2.id
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
def test_caching_reasoning_args_hit(): # test in memory cache
try:
# litellm._turn_on_debug()
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
reasoning_effort="low",
mock_response="My response",
)
response2 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
reasoning_effort="low",
mock_response="My response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response1.id == response2.id
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
def test_caching_thinking_args_miss(): # test in memory cache
try:
# litellm._turn_on_debug()
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
thinking={"type": "enabled", "budget_tokens": 1024},
mock_response="My response",
)
response2 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
mock_response="My response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response1.id != response2.id
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
def test_caching_thinking_args_hit(): # test in memory cache
try:
# litellm._turn_on_debug()
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
thinking={"type": "enabled", "budget_tokens": 1024},
mock_response="My response",
)
response2 = completion(
model="claude-3-7-sonnet-latest",
messages=messages,
caching=True,
thinking={"type": "enabled", "budget_tokens": 1024},
mock_response="My response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
assert response1.id == response2.id
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_cache_key_in_hidden_params_acompletion():
"""
Test that cache_key is present in _hidden_params on cache hits for acompletion.
Validates fix for missing x-litellm-cache-key header on proxy cache hits.
"""
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
unique_content = f"test cache key hidden params {uuid.uuid4()}"
messages = [{"role": "user", "content": unique_content}]
# First call - cache miss
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
mock_response="test response",
caching=True,
)
print(f"Response 1 _hidden_params: {response1._hidden_params}")
assert response1._hidden_params.get("cache_hit") is not True
await asyncio.sleep(0.5)
# Second call - cache hit
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
mock_response="test response",
caching=True,
)
print(f"Response 2 _hidden_params: {response2._hidden_params}")
# Verify cache hit occurred
assert response2._hidden_params.get("cache_hit") is True
# Verify cache_key is present in _hidden_params
assert "cache_key" in response2._hidden_params
assert response2._hidden_params["cache_key"] is not None
# Verify both responses have same ID (cache hit)
assert response1.id == response2.id
litellm.cache = None