ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix test_async_router_context_window_fallback

Browse Source
This commit is contained in:
Ishaan Jaff 2025-02-12 16:55:49 -08:00
parent a04204db40
commit faee508d1f
1 changed files with 13 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
tests/local_testing/test_router.py
View File

@ -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)}")