phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) undo changes that led were trying to control prisma connections

Browse source
This commit is contained in:
ishaan-jaff 2024-01-06 14:32:40 +05:30
parent 04c04d62e3
commit 6011c5c8c2
1 changed files with 9 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/proxy_server.py
View file

@ -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 global prisma_client, proxy_logging_obj, user_api_key_cache
if (
database_url is not None and prisma_client is None if database_url is not None:
): # don't re-initialize prisma client after initial init
try: try:
prisma_client = PrismaClient( prisma_client = PrismaClient(
database_url=database_url, proxy_logging_obj=proxy_logging_obj database_url=database_url, proxy_logging_obj=proxy_logging_obj
@ -320,8 +319,6 @@ async def prisma_setup(database_url: Optional[str]):
print_verbose( print_verbose(
f"Error when initializing prisma, Ensure you run pip install prisma {str(e)}" 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): def load_from_azure_key_vault(use_azure_key_vault: bool = False):
@ -535,7 +532,7 @@ class ProxyConfig:
prisma_client is not None prisma_client is not None
and litellm.get_secret("SAVE_CONFIG_TO_DB", False) == True 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 = [] _tasks = []
keys = [ keys = [
"model_list", "model_list",
@ -765,7 +762,7 @@ class ProxyConfig:
print_verbose(f"GOING INTO LITELLM.GET_SECRET!") print_verbose(f"GOING INTO LITELLM.GET_SECRET!")
database_url = litellm.get_secret(database_url) database_url = litellm.get_secret(database_url)
print_verbose(f"RETRIEVED DB URL: {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 ##
cost_tracking() cost_tracking()
### MASTER KEY ### ### MASTER KEY ###
@ -1118,6 +1115,10 @@ async def startup_event():
_run_background_health_check() _run_background_health_check()
) # start the background health check coroutine. ) # 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: if prisma_client is not None and master_key is not None:
# add master key to db # add master key to db
await generate_key_helper_fn( await generate_key_helper_fn(