(fix) test router

This commit is contained in:
ishaan-jaff 2023-12-12 09:50:44 -08:00
parent 29c579e9ca
commit f5d64a4992

View file

@ -101,6 +101,7 @@ kwargs = {"model": "azure/gpt-3.5-turbo", "messages": [{"role": "user", "content
def test_sync_fallbacks(): def test_sync_fallbacks():
try: try:
print("Test router_fallbacks: test_sync_fallbacks()")
litellm.set_verbose = True litellm.set_verbose = True
customHandler = MyCustomHandler() customHandler = MyCustomHandler()
litellm.callbacks = [customHandler] litellm.callbacks = [customHandler]
@ -112,6 +113,8 @@ def test_sync_fallbacks():
print(f"response: {response}") print(f"response: {response}")
time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 1 # 0 retries, 1 fallback assert customHandler.previous_models == 1 # 0 retries, 1 fallback
print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
router.reset() router.reset()
except Exception as e: except Exception as e:
print(e) print(e)
@ -195,6 +198,11 @@ async def test_dynamic_fallbacks_async():
Allow setting the fallback in the router.completion() call. Allow setting the fallback in the router.completion() call.
""" """
try: try:
print("Router - test_dynamic_fallbacks_async")
print("Callbacks in test_dynamic_fallbacks_async: ", litellm.callbacks)
print("Success callbacks in test_dynamic_fallbacks_async: ", litellm.success_callback)
print("Async Success callbacks in test_dynamic_fallbacks_async: ", litellm._async_success_callback)
litellm.set_verbose=True
customHandler = MyCustomHandler() customHandler = MyCustomHandler()
litellm.callbacks = [customHandler] litellm.callbacks = [customHandler]
router = Router(model_list=model_list, set_verbose=True) router = Router(model_list=model_list, set_verbose=True)
@ -203,7 +211,7 @@ async def test_dynamic_fallbacks_async():
kwargs["messages"] = [{"role": "user", "content": "Hey, how's it going?"}] kwargs["messages"] = [{"role": "user", "content": "Hey, how's it going?"}]
kwargs["fallbacks"] = [{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}] kwargs["fallbacks"] = [{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}]
response = await router.acompletion(**kwargs) response = await router.acompletion(**kwargs)
print(f"response: {response}") print(f"RESPONSE: {response}")
await asyncio.sleep(0.05) # allow a delay as success_callbacks are on a separate thread await asyncio.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 1 # 0 retries, 1 fallback assert customHandler.previous_models == 1 # 0 retries, 1 fallback
router.reset() router.reset()