From 3a8e4f31600f8dce7a420e13adfbf100e406847c Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 27 Sep 2025 09:54:54 -0700 Subject: [PATCH] test_async_embedding_azure_caching --- .../test_custom_callback_router.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index 5948fb6a2e5..1e3f5f4cfb9 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -725,16 +725,22 @@ async def test_async_embedding_azure_caching(): port=os.environ["REDIS_PORT"], password=os.environ["REDIS_PASSWORD"], ) + router = Router(model_list=[{ + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "openai/text-embedding-ada-002", + }, + }]) litellm.callbacks = [customHandler_caching] unique_time = time.time() - response1 = await litellm.aembedding( - model="azure/text-embedding-ada-002", + response1 = await router.aembedding( + model="text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, ) await asyncio.sleep(1) # set cache is async for aembedding() - response2 = await litellm.aembedding( - model="azure/text-embedding-ada-002", + response2 = await router.aembedding( + model="text-embedding-ada-002", input=[f"good morning from litellm1 {unique_time}"], caching=True, )