diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py index 4c806291ae..5f6ef2e211 100644 --- a/litellm/tests/test_router.py +++ b/litellm/tests/test_router.py @@ -139,6 +139,89 @@ def test_reading_key_from_model_list(): # test_reading_key_from_model_list() + +def test_router_azure_acompletion(): + # [PROD TEST CASE] + # This is 90% of the router use case, makes an acompletion call, acompletion + stream call and verifies it got a response + # DO NOT REMOVE THIS TEST. It's an IMP ONE. Speak to Ishaan, if you are tring to remove this + litellm.set_verbose=False + import openai + try: + print("Router Test Azure - Acompletion, Acompletion with stream") + + # remove api key from env to repro how proxy passes key to router + old_api_key = os.environ["AZURE_API_KEY"] + os.environ.pop("AZURE_API_KEY", None) + + model_list = [ + { + "model_name": "gpt-3.5-turbo", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": old_api_key, + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE") + }, + "tpm": 240000, + "rpm": 1800 + }, + { + "model_name": "gpt-3.5-turbo", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": old_api_key, + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE") + }, + "tpm": 240000, + "rpm": 1800 + } + ] + + router = Router(model_list=model_list, + routing_strategy="simple-shuffle", + set_verbose=True + ) # type: ignore + + async def test1(): + + response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "user", + "content": "hello this request will fail" + } + ] + ) + print("\n response", response) + asyncio.run(test1()) + + print("\n Testing streaming response") + async def test2(): + response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "user", + "content": "hello this request will fail" + } + ], + stream=True + ) + async for chunk in response: + if chunk is not None: + print(chunk) + asyncio.run(test2()) + print("\n Passed Streaming") + os.environ["AZURE_API_KEY"] = old_api_key + router.reset() + except Exception as e: + os.environ["AZURE_API_KEY"] = old_api_key + print(f"FAILED TEST") + pytest.fail(f"Got unexpected exception on router! - {e}") +test_router_azure_acompletion() + ### FUNCTION CALLING def test_function_calling():