fix(main.py): support ttl being set for completion, embedding, image generation calls

This commit is contained in:
Krrish Dholakia 2023-12-26 17:22:31 +05:30
parent dfd2f68c07
commit f0b6b9dce2
3 changed files with 91 additions and 1 deletions

View file

@ -466,6 +466,7 @@ def completion(
"proxy_server_request",
"preset_cache_key",
"caching_groups",
"ttl",
]
default_params = openai_params + litellm_params
non_default_params = {
@ -2158,6 +2159,7 @@ def embedding(
"model_info",
"preset_cache_key",
"caching_groups",
"ttl",
]
default_params = openai_params + litellm_params
non_default_params = {
@ -2787,6 +2789,7 @@ def image_generation(
"model_info",
"preset_cache_key",
"caching_groups",
"ttl",
]
default_params = openai_params + litellm_params
non_default_params = {

View file

@ -58,6 +58,28 @@ def test_caching_v2(): # test in memory cache
# test_caching_v2()
def test_caching_with_ttl():
try:
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
litellm.success_callback = []
litellm._async_success_callback = []
assert (
response2["choices"][0]["message"]["content"]
!= response1["choices"][0]["message"]["content"]
)
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
def test_caching_with_models_v2():
messages = [
{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}
@ -724,7 +746,7 @@ def test_get_cache_key():
pytest.fail(f"Error occurred:", e)
test_get_cache_key()
# test_get_cache_key()
# test_custom_redis_cache_params()

View file

@ -80,6 +80,71 @@ async def test_acompletion_caching_on_router():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_acompletion_caching_with_ttl_on_router():
# tests acompletion + caching on router
try:
litellm.set_verbose = True
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo-0613",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 100000,
"rpm": 10000,
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-2",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
},
"tpm": 100000,
"rpm": 10000,
},
]
messages = [
{"role": "user", "content": f"write a one sentence poem {time.time()}?"}
]
start_time = time.time()
router = Router(
model_list=model_list,
redis_host=os.environ["REDIS_HOST"],
redis_password=os.environ["REDIS_PASSWORD"],
redis_port=os.environ["REDIS_PORT"],
cache_responses=True,
timeout=30,
routing_strategy="simple-shuffle",
)
response1 = await router.acompletion(
model="gpt-3.5-turbo", messages=messages, temperature=1, ttl=0
)
print(f"response1: {response1}")
await asyncio.sleep(1) # add cache is async, async sleep for cache to get set
response2 = await router.acompletion(
model="gpt-3.5-turbo", messages=messages, temperature=1, ttl=0
)
print(f"response2: {response2}")
assert response1.id != response2.id
assert len(response1.choices[0].message.content) > 0
assert (
response1.choices[0].message.content != response2.choices[0].message.content
)
router.reset()
except litellm.Timeout as e:
end_time = time.time()
print(f"timeout error occurred: {end_time - start_time}")
pass
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_acompletion_caching_on_router_caching_groups():
# tests acompletion + caching on router