fix(proxy_server.py): fix _delete_deployment to handle base case where db_model list is empty

don't delete all router models  b/c of empty list

Fixes https://github.com/BerriAI/litellm/issues/7196
This commit is contained in:
Krrish Dholakia 2024-12-12 16:07:06 -08:00
parent 2f9fbaf421
commit f2904cbb4e
2 changed files with 76 additions and 10 deletions

View file

@ -2088,7 +2088,10 @@ class ProxyConfig:
"""
global user_config_file_path, llm_router
combined_id_list = []
if llm_router is None:
## BASE CASES ##
# if llm_router is None or db_models is empty, return 0
if llm_router is None or len(db_models) == 0:
return 0
## DB MODELS ##
@ -2418,6 +2421,19 @@ class ProxyConfig:
return config
async def _get_models_from_db(self, prisma_client: PrismaClient) -> list:
try:
new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format(
str(e)
)
)
new_models = []
return new_models
async def add_deployment(
self,
prisma_client: PrismaClient,
@ -2435,15 +2451,9 @@ class ProxyConfig:
raise ValueError(
f"Master key is not initialized or formatted. master_key={master_key}"
)
try:
new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format(
str(e)
)
)
new_models = []
new_models = await self._get_models_from_db(prisma_client=prisma_client)
# update llm router
await self._update_llm_router(
new_models=new_models, proxy_logging_obj=proxy_logging_obj

View file

@ -175,6 +175,62 @@ async def test_add_existing_deployment():
assert init_len_list == len(llm_router.model_list)
@pytest.mark.asyncio
async def test_db_error_new_model_check():
"""
- if error in db, don't delete existing models
Relevant issue: https://github.com/BerriAI/litellm/blob/ddfe687b13e9f31db2fb2322887804e3d01dd467/litellm/proxy/proxy_server.py#L2461
"""
import base64
litellm_params = LiteLLM_Params(
model="gpt-3.5-turbo",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params)
deployment_2 = Deployment(
model_name="gpt-3.5-turbo-2", litellm_params=litellm_params
)
llm_router = litellm.Router(
model_list=[
deployment.to_json(exclude_none=True),
deployment_2.to_json(exclude_none=True),
]
)
init_len_list = len(llm_router.model_list)
print(f"llm_router: {llm_router}")
master_key = "sk-1234"
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
setattr(litellm.proxy.proxy_server, "master_key", master_key)
pc = ProxyConfig()
encrypted_litellm_params = litellm_params.dict(exclude_none=True)
for k, v in encrypted_litellm_params.items():
if isinstance(v, str):
encrypted_value = encrypt_value(v, master_key)
encrypted_litellm_params[k] = base64.b64encode(encrypted_value).decode(
"utf-8"
)
db_model = DBModel(
model_id=deployment.model_info.id,
model_name="gpt-3.5-turbo",
litellm_params=encrypted_litellm_params,
model_info={"id": deployment.model_info.id},
)
db_models = []
deleted_deployments = await pc._delete_deployment(db_models=db_models)
assert deleted_deployments == 0
assert init_len_list == len(llm_router.model_list)
litellm_params = LiteLLM_Params(
model="azure/chatgpt-v-2",
api_key=os.getenv("AZURE_API_KEY"),