2024-07-04 00:48:19 +08:00
import os
import sys
2023-10-30 20:35:55 +08:00
import time
2023-08-17 12:44:50 +08:00
import traceback
2025-09-26 06:47:01 +08:00
from litellm . _uuid import uuid
2024-07-04 00:48:19 +08:00
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
2025-03-18 09:10:39 +08:00
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
2024-07-04 00:48:19 +08:00
import asyncio
import hashlib
import random
2023-08-17 12:44:50 +08:00
import pytest
2024-07-04 00:48:19 +08:00
2023-08-17 12:44:50 +08:00
import litellm
2024-07-04 00:48:19 +08:00
from litellm import aembedding , completion , embedding
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import Cache
2025-02-13 10:01:32 +08:00
from redis . asyncio import RedisCluster
from litellm . caching . redis_cluster_cache import RedisClusterCache
2024-10-22 07:42:12 +08:00
from unittest . mock import AsyncMock , patch , MagicMock , call
2024-10-07 18:23:18 +08:00
import datetime
2024-10-22 07:42:12 +08:00
from datetime import timedelta
2024-10-07 18:23:18 +08:00
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
2024-04-20 06:03:25 +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 ]
"""
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import DualCache , InMemoryCache , RedisCache
2024-04-20 06:03:25 +08:00
in_memory_cache = InMemoryCache ( )
redis_cache = RedisCache ( ) # get credentials from environment
dual_cache = DualCache ( in_memory_cache = in_memory_cache , redis_cache = redis_cache )
2024-11-04 10:18:20 +08:00
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 " }
2024-04-20 06:03:25 +08:00
2024-11-04 10:18:20 +08:00
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 " ] )
2024-04-20 06:03:25 +08:00
2024-11-04 10:18:20 +08:00
assert mock_redis_cache . call_count == 1
2024-04-20 06:03:25 +08:00
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 ]
"""
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import DualCache , InMemoryCache , RedisCache
2024-04-20 06:03:25 +08:00
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 " )
2024-10-29 12:52:12 +08:00
result = dual_cache . batch_get_cache (
keys = [ " test_value " , " test_value_2 " ] , parent_otel_span = None
)
2024-04-20 06:03:25 +08:00
assert result [ 0 ] == " hello world "
assert result [ 1 ] == None
2025-02-20 11:56:57 +08:00
@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
2024-03-18 23:21:36 +08:00
# @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
2023-12-12 10:10:46 +08:00
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 } " )
2023-12-14 18:34:45 +08:00
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
2023-12-26 19:52:31 +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 } " )
2024-03-26 04:40:17 +08:00
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 } " )
2024-04-24 07:13:03 +08:00
@pytest.mark.parametrize (
" sync_flag " ,
[ True , False ] ,
)
@pytest.mark.asyncio
async def test_caching_with_cache_controls ( sync_flag ) :
2024-01-03 15:12:30 +08:00
try :
litellm . set_verbose = True
litellm . cache = Cache ( )
message = [ { " role " : " user " , " content " : f " Hey, how ' s it going? { uuid . uuid4 ( ) } " } ]
2024-04-24 07:13:03 +08:00
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 (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " ttl " : 0 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
await asyncio . sleep ( 10 )
response2 = await litellm . acompletion (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " s-maxage " : 10 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
assert response2 [ " id " ] != response1 [ " id " ]
2024-01-03 15:12:30 +08:00
message = [ { " role " : " user " , " content " : f " Hey, how ' s it going? { uuid . uuid4 ( ) } " } ]
## TTL = 5
2024-04-24 07:13:03 +08:00
if sync_flag :
response1 = completion (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " ttl " : 5 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
response2 = completion (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " s-maxage " : 5 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
print ( f " response1: { response1 } " )
print ( f " response2: { response2 } " )
assert response2 [ " id " ] == response1 [ " id " ]
else :
response1 = await litellm . acompletion (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " ttl " : 25 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
await asyncio . sleep ( 10 )
response2 = await litellm . acompletion (
2024-08-06 02:18:59 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
cache = { " s-maxage " : 25 } ,
mock_response = " Hello world " ,
2024-04-24 07:13:03 +08:00
)
print ( f " response1: { response1 } " )
print ( f " response2: { response2 } " )
assert response2 [ " id " ] == response1 [ " id " ]
2024-01-03 15:12:30 +08:00
except Exception as e :
print ( f " error occurred: { traceback . format_exc ( ) } " )
pytest . fail ( f " Error occurred: { e } " )
# test_caching_with_cache_controls()
2025-09-10 10:48:35 +08:00
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 " )
2024-01-30 03:00:15 +08:00
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
2023-12-12 10:10:46 +08:00
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:18:59 +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
2024-08-06 02:18:59 +08:00
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
2024-08-06 02:18:59 +08:00
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
2024-01-04 14:11:23 +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 ( )
2024-01-04 14:11:23 +08:00
# 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
2023-12-12 10:10:46 +08:00
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
2025-09-10 10:48:35 +08:00
2025-04-30 12:21:28 +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
2025-09-10 10:48:35 +08:00
assert embedding3 . usage . prompt_tokens != 0
2025-04-30 12:21:28 +08:00
## 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
2025-09-10 10:48:35 +08:00
2025-05-10 14:37:02 +08:00
@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 " ]
2025-05-29 13:02:55 +08:00
assert len ( embedding3 . data ) == 1
2025-05-10 14:37:02 +08:00
assert embedding3 . _hidden_params [ " cache_hit " ] == True
2025-09-10 10:48:35 +08:00
assert embedding3 . usage . prompt_tokens != 0
2025-04-30 12:21:28 +08:00
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
2023-12-12 10:10:46 +08:00
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 2 nd 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 ( ) } "
2024-01-13 18:03:57 +08:00
embedding_1 = [ common_msg ]
2024-01-13 00:16:41 +08:00
embedding_2 = [
common_msg ,
f " I ' m fine { uuid . uuid4 ( ) } " ,
2024-01-13 18:03:57 +08:00
]
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
)
2024-01-13 18:03:57 +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
2024-04-09 03:17:57 +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 2 nd 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 " ]
` ` `
"""
2024-10-14 19:04:01 +08:00
litellm . set_verbose = True
2024-04-09 03:17:57 +08:00
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
2024-04-09 03:17:57 +08:00
)
2024-10-14 19:04:01 +08:00
print ( " embedding val 1 " , embedding_val_1 )
2024-04-09 03:17:57 +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-04-09 03:17:57 +08:00
)
2024-10-14 19:04:01 +08:00
print ( " embedding val 2 " , embedding_val_2 )
2024-04-09 03:17:57 +08:00
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
2024-08-28 13:52:11 +08:00
@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 " ] ,
)
2025-09-26 06:47:01 +08:00
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 \n CALL2 \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 " ]
2024-10-07 18:23:18 +08:00
@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 " ,
2024-10-30 13:04:16 +08:00
default_in_redis_ttl = 2 ,
2024-10-07 18:23:18 +08:00
)
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 " ,
2024-10-07 18:23:18 +08:00
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 (
2024-10-30 13:04:16 +08:00
seconds = 2
2024-10-07 18:23:18 +08:00
) # 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
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
2025-09-26 06:47:01 +08:00
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 ( ) :
2023-11-24 00:27:44 +08:00
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 " ] ,
)
2024-01-05 18:33:56 +08:00
print ( " test2 for Redis Caching - non streaming " )
2023-12-25 16:40:38 +08:00
response1 = completion (
2024-05-10 16:57:24 +08:00
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 )
2023-11-24 00:27:44 +08:00
print ( " \n response 1 " , response1 )
print ( " \n response 2 " , response2 )
print ( " \n response 3 " , response3 )
2023-11-24 10:17:11 +08:00
print ( " \n response 4 " , response4 )
2023-08-29 13:10:15 +08:00
litellm . cache = None
2023-12-12 10:10:46 +08:00
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
2024-01-05 18:33:56 +08:00
assert response1 . id == response2 . id
assert response1 . created == response2 . created
assert response1 . choices [ 0 ] . message . content == response2 . choices [ 0 ] . message . content
2024-07-04 00:48:19 +08:00
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 = [ ]
2023-12-10 03:32:14 +08:00
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 ,
)
2024-03-06 13:12:50 +08:00
response_1_id = " "
2023-12-09 03:50:36 +08:00
for chunk in response1 :
print ( chunk )
2024-03-06 13:12:50 +08:00
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 ,
)
2024-03-06 13:12:50 +08:00
response_2_id = " "
2023-12-09 03:50:36 +08:00
for chunk in response2 :
print ( chunk )
2024-03-26 12:36:47 +08:00
response_2_id = chunk . id
2023-12-25 16:40:38 +08:00
assert (
2024-03-06 13:12:50 +08:00
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
2023-12-12 10:10:46 +08:00
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. " )
2024-08-22 06:01:52 +08:00
@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
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import RedisCache
2024-08-22 06:01:52 +08:00
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 )
2024-09-07 23:56:00 +08:00
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
2024-09-08 00:37:23 +08:00
@pytest.mark.skip ( reason = " Local test. Requires running redis cluster locally. " )
2024-09-07 23:56:00 +08:00
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
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import RedisCache
2024-09-07 23:56:00 +08:00
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 )
2024-09-08 00:37:23 +08:00
print ( " REDIS_CLUSTER_NODES " , os . environ [ " REDIS_CLUSTER_NODES " ] )
2024-09-07 23:56:00 +08:00
# 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 )
2024-08-22 06:01:52 +08:00
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 )
2024-10-19 17:51:11 +08:00
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 ( " \n response 1 " , response_1_content )
print ( " \n 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-10 05:56:53 +08:00
litellm . cache = None
2023-12-12 10:10:46 +08:00
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
2023-12-12 10:10:46 +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 } " ,
}
]
2024-03-18 23:21:36 +08:00
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 (
2025-09-10 10:48:35 +08:00
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 )
2024-10-16 15:03:49 +08:00
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 (
2025-09-10 10:48:35 +08:00
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 )
2024-10-16 15:03:49 +08:00
print ( " \n final response 1 " , response_1_content )
print ( " \n final 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 } "
2024-01-05 13:21:56 +08:00
2023-12-12 00:43:00 +08:00
litellm . cache = None
2023-12-12 10:10:46 +08:00
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
2024-07-04 00:48:19 +08:00
2024-08-20 01:56:04 +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
2024-08-20 01:56:04 +08:00
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 } " ,
}
]
2024-01-03 18:12:09 +08:00
litellm . cache = Cache (
2024-03-08 00:56:59 +08:00
type = " s3 " ,
2024-08-20 01:56:04 +08:00
s3_bucket_name = " litellm-proxy " ,
s3_region_name = " us-west-2 " ,
2024-01-03 18:12:09 +08:00
)
print ( " s3 Cache: test for caching, streaming + completion " )
2024-01-03 17:43:39 +08:00
response_1_content = " "
response_2_content = " "
2024-01-05 18:33:56 +08:00
response_1_created = " "
response_2_created = " "
2024-08-20 01:56:04 +08:00
if sync_mode :
response1 = litellm . completion (
2025-10-26 01:19:24 +08:00
model = " azure/gpt-4.1-mini " ,
2024-08-20 01:56:04 +08:00
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 " ,
2024-08-20 01:56:04 +08:00
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 " )
2024-08-20 01:56:04 +08:00
if sync_mode :
response2 = litellm . completion (
2025-10-26 01:19:24 +08:00
model = " azure/gpt-4.1-mini " ,
2024-08-20 01:56:04 +08:00
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 " ,
2024-08-20 01:56:04 +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 " "
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 ( " \n response 1 " , response_1_content )
print ( " \n response 2 " , response_2_content )
2024-01-05 18:33:56 +08:00
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 18:33:56 +08:00
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 18:33:56 +08:00
2024-01-05 19:10:56 +08:00
# assert response_1_created == response_2_created
2024-01-05 18:33:56 +08:00
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
2024-01-05 18:33:56 +08:00
# test_s3_cache_acompletion_stream_azure()
2024-01-03 17:43:39 +08:00
2024-04-05 06:07:19 +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 (
2024-03-07 10:07:28 +08:00
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
2023-12-12 10:10:46 +08:00
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
2023-12-14 18:50:29 +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
2023-12-14 18:50:29 +08:00
litellm . cache = Cache ( )
print ( " Testing cache override " )
2023-12-25 16:40:38 +08:00
litellm . set_verbose = True
2023-12-14 18:50:29 +08:00
# test embedding
response1 = embedding (
2023-12-25 16:40:38 +08:00
model = " text-embedding-ada-002 " , input = [ " hello who are you " ] , caching = False
2023-12-14 18:50:29 +08:00
)
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
2023-12-14 18:50:29 +08:00
)
end_time = time . time ( )
print ( f " Embedding 2 response time: { end_time - start_time } seconds " )
2023-12-25 16:40:38 +08:00
assert (
2024-01-27 15:00:59 +08:00
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-14 18:50:29 +08:00
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
2025-09-26 06:47:01 +08:00
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 ,
}
] ,
2024-04-24 03:58:30 +08:00
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 ,
}
] ,
2024-04-24 03:58:30 +08:00
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
2025-09-26 06:47:01 +08:00
from litellm . _uuid import uuid
2024-04-24 03:58:30 +08:00
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
2023-12-12 10:10:46 +08:00
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 ( ) :
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import Cache
2023-12-25 16:40:38 +08:00
2023-12-10 08:42:54 +08:00
try :
2023-12-14 16:38:58 +08:00
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 " : { } ,
}
2023-12-14 16:38:58 +08:00
)
2024-08-06 02:18:59 +08:00
cache_key_str = " model: gpt-3.5-turbomessages: [ { ' role ' : ' user ' , ' content ' : ' write a one sentence poem about: 7510 ' }]max_tokens: 40temperature: 0.2stream: True "
2023-12-29 17:33:33 +08:00
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 (
2023-12-29 17:33:33 +08:00
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 "
2023-12-14 16:38:58 +08:00
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> " ,
2023-12-14 16:38:58 +08:00
}
)
print ( embedding_cache_key )
2023-12-29 17:33:33 +08:00
embedding_cache_key_str = (
2025-09-28 03:41:35 +08:00
" model: azure/text-embedding-ada-002input: [ ' hi who is ishaan ' ] "
2023-12-29 17:33:33 +08:00
)
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 (
2023-12-29 17:33:33 +08:00
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 "
2023-12-14 16:38:58 +08:00
# 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> " ,
}
2023-12-14 16:38:58 +08:00
)
print ( embedding_cache_key_2 )
2023-12-29 17:33:33 +08:00
embedding_cache_key_str_2 = (
" model: EMBEDDING_MODEL_GROUPinput: [ ' hi who is ishaan ' ] "
2023-12-25 16:40:38 +08:00
)
2023-12-29 17:33:33 +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
2023-12-14 16:38:58 +08:00
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
2023-12-26 19:52:31 +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
2024-07-04 00:48:19 +08:00
2024-07-18 02:15:30 +08:00
def test_caching_redis_simple ( caplog , capsys ) :
2024-07-04 00:48:19 +08:00
"""
Relevant issue - https : / / github . com / BerriAI / litellm / issues / 4511
"""
2024-07-18 02:15:30 +08:00
litellm . set_verbose = True ## REQUIRED FOR TEST.
2024-07-04 00:48:19 +08:00
litellm . cache = Cache (
type = " redis " , url = os . getenv ( " REDIS_SSL_URL " )
) # passing `supported_call_types = ["completion"]` has no effect
s = time . time ( )
2024-07-18 02:15:30 +08:00
uuid_str = str ( uuid . uuid4 ( ) )
2024-07-04 00:48:19 +08:00
x = completion (
2024-07-18 02:15:30 +08:00
model = " gpt-3.5-turbo " ,
messages = [ { " role " : " user " , " content " : f " Hello, how are you? Wink { uuid_str } " } ] ,
2024-07-04 00:48:19 +08:00
stream = True ,
)
for m in x :
print ( m )
print ( time . time ( ) - s )
s2 = time . time ( )
x = completion (
2024-07-18 02:15:30 +08:00
model = " gpt-3.5-turbo " ,
messages = [ { " role " : " user " , " content " : f " Hello, how are you? Wink { uuid_str } " } ] ,
2024-07-04 00:48:19 +08:00
stream = True ,
)
for m in x :
print ( m )
print ( time . time ( ) - s2 )
2024-07-07 00:14:29 +08:00
redis_async_caching_error = False
redis_service_logging_error = False
2024-07-18 02:15:30 +08:00
captured = capsys . readouterr ( )
2024-07-04 00:48:19 +08:00
captured_logs = [ rec . message for rec in caplog . records ]
2024-07-07 00:14:29 +08:00
print ( f " captured_logs: { captured_logs } " )
for item in captured_logs :
2024-07-18 02:15:30 +08:00
if (
" Error connecting to Async Redis client " in item
or " Set ASYNC Redis Cache " in item
) :
2024-07-07 00:14:29 +08:00
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
2024-07-18 02:15:30 +08:00
assert " async success_callback: reaches cache for logging " not in captured . out
2024-08-20 02:59:50 +08:00
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 )
2024-10-14 19:04:01 +08:00
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
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
2024-09-29 04:00:45 +08:00
@pytest.mark.asyncio
2025-07-20 06:58:25 +08:00
@pytest.mark.flaky ( retries = 3 , delay = 2 )
2024-09-29 04:00:45 +08:00
async def test_redis_proxy_batch_redis_get_cache ( ) :
"""
Tests batch_redis_get . py
- make 1 st call - > expect miss
- make 2 nd call - > expect hit
"""
2024-10-14 19:04:01 +08:00
from litellm . caching . caching import Cache , DualCache
2024-09-29 04:00:45 +08:00
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 ( )
2025-09-26 06:47:01 +08:00
from litellm . _uuid import uuid
2024-09-29 04:00:45 +08:00
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
2024-10-10 15:42:11 +08:00
2025-03-18 09:10:39 +08:00
@pytest.mark.parametrize ( " sync_mode " , [ True , False ] )
@pytest.mark.asyncio
async def test_logging_turn_off_message_logging_streaming ( sync_mode ) :
2024-10-10 15:42:11 +08:00
litellm . turn_off_message_logging = True
mock_obj = Cache ( type = " local " )
litellm . cache = mock_obj
2025-03-18 09:10:39 +08:00
with patch . object ( mock_obj , " add_cache " ) as mock_client , patch . object (
mock_obj , " async_add_cache "
) as mock_async_client :
2024-10-10 15:42:11 +08:00
print ( f " mock_obj.add_cache: { mock_obj . add_cache } " )
2025-03-18 09:10:39 +08:00
if sync_mode is True :
resp = litellm . completion (
model = " gpt-3.5-turbo " ,
messages = [ { " role " : " user " , " content " : " hi " } ] ,
mock_response = " hello " ,
stream = True ,
)
2024-10-10 15:42:11 +08:00
2025-03-18 09:10:39 +08:00
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 ,
)
2024-10-10 15:42:11 +08:00
2025-03-18 09:10:39 +08:00
async for chunk in resp :
continue
2024-10-10 15:42:11 +08:00
2025-03-18 09:10:39 +08:00
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
2024-10-13 02:48:34 +08:00
2024-10-15 11:17:23 +08:00
def test_basic_caching_import ( ) :
from litellm . caching import Cache
assert Cache is not None
print ( " Cache imported successfully " )
2024-10-17 13:16:23 +08:00
@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 ,
2025-08-04 23:36:49 +08:00
CompletionTokensDetailsWrapper ,
PromptTokensDetailsWrapper ,
2024-10-17 13:16:23 +08:00
)
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 ,
2025-08-04 23:36:49 +08:00
completion_tokens_details = CompletionTokensDetailsWrapper (
2024-10-17 13:16:23 +08:00
audio_tokens = None , reasoning_tokens = 0
) ,
2025-08-04 23:36:49 +08:00
prompt_tokens_details = PromptTokensDetailsWrapper (
2024-10-17 13:16:23 +08:00
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 )
2024-10-20 07:16:51 +08:00
@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
2024-10-22 07:42:12 +08:00
def test_redis_caching_default_ttl ( ) :
"""
Ensure that the default redis cache TTL is 60 s
"""
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
2024-10-22 07:42:12 +08:00
)
## 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 (
[
2025-02-13 10:01:32 +08:00
call . set (
name = " test_key1 " , value = ' " test_value1 " ' , ex = expected_timedelta
) ,
call . set (
name = " test_key2 " , value = ' " test_value2 " ' , ex = expected_timedelta
) ,
2024-10-22 07:42:12 +08:00
]
)
@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
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
2024-11-25 08:36:19 +08:00
@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
2025-02-13 10:01:32 +08:00
litellm . _turn_on_debug ( )
2024-11-25 08:36:19 +08:00
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
2025-09-10 10:48:35 +08:00
assert len ( results ) == 4
2024-11-25 08:36:19 +08:00
# 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
2024-12-07 13:14:12 +08:00
@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
2025-02-19 06:47:34 +08:00
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
"""
2025-09-26 06:47:01 +08:00
from litellm . _uuid import uuid
2025-07-13 03:06:16 +08:00
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 } " } ]
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# Mock the Redis client creation from the _redis module
2025-09-10 10:48:35 +08:00
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 :
2025-07-13 03:06:16 +08:00
# Create a mock Redis client that simulates real Redis behavior
mock_redis_client = MagicMock ( )
mock_get_redis_client . return_value = mock_redis_client
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# Mock the connection pool
mock_connection_pool = MagicMock ( )
mock_get_redis_connection_pool . return_value = mock_connection_pool
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# Dictionary to simulate Redis storage with namespace support
redis_storage = { }
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
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
2025-09-10 10:48:35 +08:00
return json . dumps ( value ) . encode ( " utf-8 " )
2025-07-13 03:06:16 +08:00
return None
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
def mock_redis_set ( name , value , ex = None , * * kwargs ) :
print ( f " Redis SET: { name } = { value } " )
redis_storage [ name ] = value
return True
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
def mock_redis_ping ( ) :
return True
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
def mock_redis_info ( ) :
return { " redis_version " : " 7.0.0 " }
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
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
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# Initialize the cache
litellm . cache = Cache ( type = " redis " )
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
namespace_1 = " org-id1 "
namespace_2 = " org-id2 "
# Use mock_response to ensure deterministic responses without external API calls
response_1 = completion (
2025-09-10 10:48:35 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
2025-07-13 03:06:16 +08:00
cache = { " namespace " : namespace_1 } ,
2025-09-10 10:48:35 +08:00
mock_response = " Response for namespace 1 " ,
2025-07-13 03:06:16 +08:00
)
2025-02-19 06:47:34 +08:00
2025-07-13 03:06:16 +08:00
response_2 = completion (
2025-09-10 10:48:35 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
2025-07-13 03:06:16 +08:00
cache = { " namespace " : namespace_2 } ,
2025-09-10 10:48:35 +08:00
mock_response = " Response for namespace 2 " ,
2025-07-13 03:06:16 +08:00
)
2025-02-19 06:47:34 +08:00
2025-07-13 03:06:16 +08:00
response_3 = completion (
2025-09-10 10:48:35 +08:00
model = " gpt-3.5-turbo " ,
messages = messages ,
2025-07-13 03:06:16 +08:00
cache = { " namespace " : namespace_1 } ,
2025-09-10 10:48:35 +08:00
mock_response = " This should be cached " ,
2025-07-13 03:06:16 +08:00
)
2025-02-19 06:47:34 +08:00
2025-07-13 03:06:16 +08:00
response_4 = completion (
2025-09-10 10:48:35 +08:00
model = " gpt-3.5-turbo " ,
2025-07-13 03:06:16 +08:00
messages = messages ,
2025-09-10 10:48:35 +08:00
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 ' ) } "
2025-07-13 03:06:16 +08:00
)
2025-02-19 06:47:34 +08:00
2025-07-13 03:06:16 +08:00
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 } : " ) ]
2025-09-10 10:48:35 +08:00
no_namespace_keys = [
k
for k in cache_keys
if not k . startswith ( f " { namespace_1 } : " )
and not k . startswith ( f " { namespace_2 } : " )
]
2025-07-13 03:06:16 +08:00
print ( f " Namespace 1 keys: { namespace_1_keys } " )
print ( f " Namespace 2 keys: { namespace_2_keys } " )
print ( f " No namespace keys: { no_namespace_keys } " )
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# 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 "
2025-09-10 10:48:35 +08:00
2025-07-13 03:06:16 +08:00
# The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1
2025-09-10 10:48:35 +08:00
assert isinstance (
response_3 , str
) , " Response 3 should be a cache hit (string) for same namespace "
2025-07-13 03:06:16 +08:00
# response 1 & 2 should be ModelResponse objects (cache misses)
2025-09-10 10:48:35 +08:00
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 "
2025-07-13 03:06:16 +08:00
# response 1 & 2 should have different IDs (different namespaces)
2025-09-10 10:48:35 +08:00
assert (
response_1 . id != response_2 . id
) , f " Expected different response ID for different namespace. Got { response_1 . id } and { response_2 . id } "
2025-07-13 03:06:16 +08:00
# response 1 & 4 should have different IDs (different namespaces)
2025-09-10 10:48:35 +08:00
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 } "
2025-03-05 13:12:16 +08:00
2025-12-07 01:41:09 +08:00
@pytest.mark.flaky ( retries = 3 , delay = 1 )
2025-03-05 13:12:16 +08:00
def test_caching_with_reasoning_content ( ) :
"""
Test that reasoning content is cached
"""
2025-09-26 06:47:01 +08:00
from litellm . _uuid import uuid
2025-03-05 13:12:16 +08:00
2025-12-07 01:41:09 +08:00
try :
messages = [ { " role " : " user " , " content " : f " what is litellm? { uuid . uuid4 ( ) } " } ]
litellm . cache = Cache ( )
2025-03-05 13:12:16 +08:00
2025-12-07 01:41:09 +08:00
response_1 = completion (
model = " anthropic/claude-3-7-sonnet-latest " ,
messages = messages ,
thinking = { " type " : " enabled " , " budget_tokens " : 1024 } ,
)
2025-03-05 13:12:16 +08:00
2025-12-07 01:41:09 +08:00
response_2 = completion (
model = " anthropic/claude-3-7-sonnet-latest " ,
messages = messages ,
thinking = { " type " : " enabled " , " budget_tokens " : 1024 } ,
)
2025-03-05 13:12:16 +08:00
2025-12-07 01:41:09 +08:00
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 ) } " )
2025-04-22 13:39:40 +08:00
def test_caching_reasoning_args_miss ( ) : # test in memory cache
try :
2025-09-10 10:48:35 +08:00
# litellm._turn_on_debug()
2025-04-22 13:39:40 +08:00
litellm . set_verbose = True
2025-09-10 10:48:35 +08:00
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 " ,
2025-04-22 13:39:40 +08:00
)
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 } " )
2025-09-10 10:48:35 +08:00
2025-04-22 13:39:40 +08:00
def test_caching_reasoning_args_hit ( ) : # test in memory cache
try :
2025-09-10 10:48:35 +08:00
# litellm._turn_on_debug()
2025-04-22 13:39:40 +08:00
litellm . set_verbose = True
2025-09-10 10:48:35 +08:00
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 " ,
2025-04-22 13:39:40 +08:00
)
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 } " )
2025-09-10 10:48:35 +08:00
2025-04-22 13:39:40 +08:00
def test_caching_thinking_args_miss ( ) : # test in memory cache
try :
2025-09-10 10:48:35 +08:00
# litellm._turn_on_debug()
2025-04-22 13:39:40 +08:00
litellm . set_verbose = True
2025-09-10 10:48:35 +08:00
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 " ,
2025-04-22 13:39:40 +08:00
)
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 } " )
2025-09-10 10:48:35 +08:00
2025-04-22 13:39:40 +08:00
def test_caching_thinking_args_hit ( ) : # test in memory cache
try :
2025-09-10 10:48:35 +08:00
# litellm._turn_on_debug()
2025-04-22 13:39:40 +08:00
litellm . set_verbose = True
2025-09-10 10:48:35 +08:00
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 " ,
2025-04-22 13:39:40 +08:00
)
print ( f " response1: { response1 } " )
print ( f " response2: { response2 } " )
assert response1 . id == response2 . id
except Exception as e :
print ( f " error occurred: { traceback . format_exc ( ) } " )
2025-04-30 12:21:28 +08:00
pytest . fail ( f " Error occurred: { e } " )
2025-10-09 09:10:43 +08:00
@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