diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index ee51c15101..f4fbeaff1c 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -718,61 +718,9 @@ def test_router_azure_acompletion(): pytest.fail(f"Got unexpected exception on router! - {e}") -# test_router_azure_acompletion() - - -def test_router_context_window_fallback(): - """ - - Give a gpt-3.5-turbo model group with different context windows (4k vs. 16k) - - Send a 5k prompt - - Assert it works - """ - import os - - from large_text import text - - litellm.set_verbose = False - - print(f"len(text): {len(text)}") - try: - 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": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), - "base_model": "azure/gpt-35-turbo", - }, - }, - { - "model_name": "gpt-3.5-turbo-large", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "gpt-3.5-turbo-1106", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - ] - - router = Router(model_list=model_list, set_verbose=True, context_window_fallbacks=[{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}], num_retries=0) # type: ignore - - response = router.completion( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": text}, - {"role": "user", "content": "Who was Alexander?"}, - ], - ) - - print(f"response: {response}") - assert response.model == "gpt-3.5-turbo-1106" - except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") - - @pytest.mark.asyncio -async def test_async_router_context_window_fallback(): +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_async_router_context_window_fallback(sync_mode): """ - Give a gpt-4 model group with different context windows (8192k vs. 128k) - Send a 10k prompt @@ -806,17 +754,23 @@ async def test_async_router_context_window_fallback(): ] router = Router(model_list=model_list, set_verbose=True, context_window_fallbacks=[{"gpt-4": ["gpt-4-turbo"]}], num_retries=0) # type: ignore + if sync_mode is False: + response = await router.acompletion( + model="gpt-4", + messages=[ + {"role": "system", "content": text * 2}, + {"role": "user", "content": "Who was Alexander?"}, + ], + ) - response = await router.acompletion( - model="gpt-4", - messages=[ - {"role": "system", "content": text * 2}, - {"role": "user", "content": "Who was Alexander?"}, - ], - ) - - print(f"response: {response}") - assert "gpt-4-turbo" in response.model + print(f"response: {response}") + assert "gpt-4-turbo" in response.model + else: + response = router.completion( + model="gpt-4", + messages=[{"role": "user", "content": "Who was Alexander?"}], + ) + assert "gpt-4-turbo" in response.model except Exception as e: pytest.fail(f"Got unexpected exception on router! - {str(e)}")