test(test_models.py): fix delete model test

This commit is contained in:
Krrish Dholakia 2024-04-04 08:45:32 -07:00
parent 674499b41a
commit 4b56f08cbe
3 changed files with 31 additions and 26 deletions

View file

@ -215,7 +215,7 @@ class ProxyChatCompletionRequest(LiteLLMBase):
class ModelInfoDelete(LiteLLMBase):
id: Optional[str]
id: str
class ModelInfo(LiteLLMBase):

View file

@ -6833,6 +6833,7 @@ async def add_new_model(
description="v2 - returns all the models set on the config.yaml, shows 'user_access' = True if the user has access to the model. Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
tags=["model management"],
dependencies=[Depends(user_api_key_auth)],
include_in_schema=False,
)
async def model_info_v2(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -7036,37 +7037,40 @@ async def model_info_v1(
async def delete_model(model_info: ModelInfoDelete):
global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config
try:
if not os.path.exists(user_config_file_path):
raise HTTPException(status_code=404, detail="Config file does not exist.")
"""
[BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
# Load existing config
config = await proxy_config.get_config()
- Check if id in db
- Delete
"""
# If model_list is not in the config, nothing can be deleted
if len(config.get("model_list", [])) == 0:
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400, detail="No model list available in the config."
status_code=500,
detail={
"error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
},
)
# Check if the model with the specified model_id exists
model_to_delete = None
for model in config["model_list"]:
if model.get("model_info", {}).get("id", None) == model_info.id:
model_to_delete = model
break
# If the model was not found, return an error
if model_to_delete is None:
raise HTTPException(
status_code=400, detail="Model with given model_id not found."
# update DB
if general_settings.get("store_model_in_db", False) == True:
"""
- store model_list in db
- store keys separately
"""
# encrypt litellm params #
await prisma_client.db.litellm_proxymodeltable.delete(
where={"model_id": model_info.id}
)
else:
raise HTTPException(
status_code=500,
detail={
"error": "Set `store_model_in_db: true` in general_settings on your config.yaml"
},
)
# Remove model from the list and save the updated config
config["model_list"].remove(model_to_delete)
# Save updated config
config = await proxy_config.save_config(new_config=config)
return {"message": "Model deleted successfully"}
except Exception as e:

View file

@ -186,5 +186,6 @@ async def test_delete_models():
key_gen = await generate_key(session=session)
key = key_gen["key"]
await add_models(session=session, model_id=model_id)
await asyncio.sleep(60)
await chat_completion(session=session, key=key)
await delete_model(session=session, model_id=model_id)