From faee508d1f646343d1c516af74ead10753351f47 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 12 Feb 2025 16:55:49 -0800 Subject: [PATCH] fix test_async_router_context_window_fallback --- tests/local_testing/test_router.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 62d0a5f52e..ee51c15101 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -774,8 +774,8 @@ def test_router_context_window_fallback(): @pytest.mark.asyncio async def test_async_router_context_window_fallback(): """ - - Give a gpt-3.5-turbo model group with different context windows (4k vs. 16k) - - Send a 5k prompt + - Give a gpt-4 model group with different context windows (8192k vs. 128k) + - Send a 10k prompt - Assert it works """ import os @@ -783,41 +783,40 @@ async def test_async_router_context_window_fallback(): from large_text import text litellm.set_verbose = False + litellm._turn_on_debug() print(f"len(text): {len(text)}") try: model_list = [ { - "model_name": "gpt-3.5-turbo", # openai model name + "model_name": "gpt-4", # 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": "gpt-4", + "api_key": os.getenv("OPENAI_API_KEY"), + "api_base": os.getenv("OPENAI_API_BASE"), }, }, { - "model_name": "gpt-3.5-turbo-large", # openai model name + "model_name": "gpt-4-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call - "model": "gpt-3.5-turbo-1106", + "model": "gpt-4-turbo", "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 + router = Router(model_list=model_list, set_verbose=True, context_window_fallbacks=[{"gpt-4": ["gpt-4-turbo"]}], num_retries=0) # type: ignore response = await router.acompletion( - model="gpt-3.5-turbo", + model="gpt-4", messages=[ - {"role": "system", "content": text}, + {"role": "system", "content": text * 2}, {"role": "user", "content": "Who was Alexander?"}, ], ) print(f"response: {response}") - assert response.model == "gpt-3.5-turbo-1106" + assert "gpt-4-turbo" in response.model except Exception as e: pytest.fail(f"Got unexpected exception on router! - {str(e)}")