From 4b56f08cbe69a0f22b646bfb584e3c0a22bcb83c Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 4 Apr 2024 08:45:32 -0700
Subject: [PATCH] test(test_models.py): fix delete model test

---
 litellm/proxy/_types.py       |  2 +-
 litellm/proxy/proxy_server.py | 54 +++++++++++++++++++----------------
 tests/test_models.py          |  1 +
 3 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 19107bb3d..e809d144d 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -215,7 +215,7 @@ class ProxyChatCompletionRequest(LiteLLMBase):
 
 
 class ModelInfoDelete(LiteLLMBase):
-    id: Optional[str]
+    id: str
 
 
 class ModelInfo(LiteLLMBase):
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 2e9616821..a69f28d10 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -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:
diff --git a/tests/test_models.py b/tests/test_models.py
index b76dfb116..5b4821529 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -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)