From 6011c5c8c2fae6d2b8e6145deffdbb2822eb9d06 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 6 Jan 2024 14:32:40 +0530 Subject: [PATCH] (fix) undo changes that led were trying to control prisma connections --- litellm/proxy/proxy_server.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f12fe5b6b0..56fc298aeb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -307,11 +307,10 @@ async def user_api_key_auth( ) -async def prisma_setup(database_url: Optional[str]): +def prisma_setup(database_url: Optional[str]): global prisma_client, proxy_logging_obj, user_api_key_cache - if ( - database_url is not None and prisma_client is None - ): # don't re-initialize prisma client after initial init + + if database_url is not None: try: prisma_client = PrismaClient( database_url=database_url, proxy_logging_obj=proxy_logging_obj @@ -320,8 +319,6 @@ async def prisma_setup(database_url: Optional[str]): print_verbose( f"Error when initializing prisma, Ensure you run pip install prisma {str(e)}" ) - if prisma_client is not None and prisma_client.db.is_connected() == False: - await prisma_client.connect() def load_from_azure_key_vault(use_azure_key_vault: bool = False): @@ -535,7 +532,7 @@ class ProxyConfig: prisma_client is not None and litellm.get_secret("SAVE_CONFIG_TO_DB", False) == True ): - await prisma_setup(database_url=None) # in case it's not been connected yet + prisma_setup(database_url=None) # in case it's not been connected yet _tasks = [] keys = [ "model_list", @@ -765,7 +762,7 @@ class ProxyConfig: print_verbose(f"GOING INTO LITELLM.GET_SECRET!") database_url = litellm.get_secret(database_url) print_verbose(f"RETRIEVED DB URL: {database_url}") - await prisma_setup(database_url=database_url) + prisma_setup(database_url=database_url) ## COST TRACKING ## cost_tracking() ### MASTER KEY ### @@ -1118,6 +1115,10 @@ async def startup_event(): _run_background_health_check() ) # start the background health check coroutine. + print_verbose(f"prisma client - {prisma_client}") + if prisma_client: + await prisma_client.connect() + if prisma_client is not None and master_key is not None: # add master key to db await generate_key_helper_fn(