From 3091d73936ea71dd65bb70b4cf3bf0c6413ee05c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 24 Jul 2025 14:15:32 -0700 Subject: [PATCH] update test auto_router_embedding_model --- .../local_testing/test_router_auto_router.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_router_auto_router.py b/tests/local_testing/test_router_auto_router.py index d4d9823da16..71d2e2f7378 100644 --- a/tests/local_testing/test_router_auto_router.py +++ b/tests/local_testing/test_router_auto_router.py @@ -14,6 +14,20 @@ from litellm import Router router = Router( model_list=[ + { + "model_name": "custom-text-embedding-model", + "litellm_params": { + "model": "text-embedding-3-large", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "custom-text-embedding-model-2", + "litellm_params": { + "model": "text-embedding-3-large", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, { "model_name": "gpt-4.1", "litellm_params": { @@ -33,14 +47,18 @@ router = Router( "model_name": "auto_router1", "litellm_params": { "model": "auto_router/auto_router_1", - "router_config_path": "auto_router/router.json", + "auto_router_config_path": "auto_router/router.json", + "auto_router_default_model": "openai/gpt-4o-mini", + "auto_router_embedding_model": "custom-text-embedding-model", }, }, { "model_name": "auto_router_2", "litellm_params": { "model": "auto_router/auto_router_2", - "router_config_path": "auto_router/router_2.json", + "auto_router_config_path": "auto_router/router_2.json", + "auto_router_default_model": "openai/gpt-4o-mini", + "auto_router_embedding_model": "custom-text-embedding-model-2", }, }, ],