From ffb752d8787610ace639b4d284ab892b51bdfb76 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 26 Aug 2024 20:32:23 -0700
Subject: [PATCH] fix entrypoint

---
 litellm/main.py                   | 3 +--
 litellm/proxy/prisma_migration.py | 4 ++++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index dd6a9e1b65..f8f036d1f6 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -82,7 +82,6 @@ from .llms import (
     bedrock,
     clarifai,
     cloudflare,
-    cohere,
     gemini,
     huggingface_restapi,
     maritalk,
@@ -107,7 +106,7 @@ from .llms.azure import AzureChatCompletion, _check_dynamic_azure_params
 from .llms.azure_text import AzureTextCompletion
 from .llms.bedrock_httpx import BedrockConverseLLM, BedrockLLM
 from .llms.cohere import chat as cohere_chat
-from .llms.cohere import completion as cohere_completion
+from .llms.cohere import completion as cohere_completion  # type: ignore
 from .llms.cohere import embed as cohere_embed
 from .llms.custom_llm import CustomLLM, custom_chat_llm_router
 from .llms.databricks import DatabricksChatCompletion
diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py
index 6ee09c22b6..c79649f6f2 100644
--- a/litellm/proxy/prisma_migration.py
+++ b/litellm/proxy/prisma_migration.py
@@ -51,6 +51,10 @@ while retry_count < max_retries and exit_code != 0:
     retry_count += 1
     print(f"Attempt {retry_count}...")  # noqa
 
+    # run prisma generate
+    result = subprocess.run(["prisma", "generate"], capture_output=True)
+    exit_code = result.returncode
+
     # Run the Prisma db push command
     result = subprocess.run(
         ["prisma", "db", "push", "--accept-data-loss"], capture_output=True