phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 19:54:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(test) test_custom_callback - aembedding cache

Browse source
This commit is contained in:
ishaan-jaff 2023-12-14 17:36:29 +05:30
parent 773b1adf5e
commit 339da65fae
1 changed files with 7 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

7
litellm/tests/test_custom_callback_input.py
View file

@ -609,6 +609,7 @@ async def test_async_completion_azure_caching():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_embedding_azure_caching(): async def test_async_embedding_azure_caching():
print("Testing custom callback input - Azure Caching")
customHandler_caching = CompletionCustomHandler() customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD']) litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
litellm.callbacks = [customHandler_caching] litellm.callbacks = [customHandler_caching]
@ -616,9 +617,15 @@ async def test_async_embedding_azure_caching():
response1 = await litellm.aembedding(model="azure/azure-embedding-model", response1 = await litellm.aembedding(model="azure/azure-embedding-model",
input=[f"good morning from litellm1 {unique_time}"], input=[f"good morning from litellm1 {unique_time}"],
caching=True) caching=True)
await asyncio.sleep(1) # set cache is async for aembedding()
response2 = await litellm.aembedding(model="azure/azure-embedding-model", response2 = await litellm.aembedding(model="azure/azure-embedding-model",
input=[f"good morning from litellm1 {unique_time}"], input=[f"good morning from litellm1 {unique_time}"],
caching=True) caching=True)
await asyncio.sleep(1) # success callbacks are done in parallel await asyncio.sleep(1) # success callbacks are done in parallel
print(customHandler_caching.states)
assert len(customHandler_caching.errors) == 0 assert len(customHandler_caching.errors) == 0
assert len(customHandler_caching.states) == 4 # pre, post, success, success assert len(customHandler_caching.states) == 4 # pre, post, success, success
# asyncio.run(
# test_async_embedding_azure_caching()
# )