mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(main.py): support ttl being set for completion, embedding, image generation calls
This commit is contained in:
parent
dfd2f68c07
commit
f0b6b9dce2
3 changed files with 91 additions and 1 deletions
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue