phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy_server.py): don't reconnect prisma if already connected

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-09 11:38:42 +05:30
parent e97eff4243
commit cd350ab8d8
3 changed files with 8 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/proxy_server.py
View file

@ -1156,7 +1156,7 @@ async def startup_event():
### CONNECT TO DB ### ### CONNECT TO DB ###
# check if DATABASE_URL in environment - load from there # check if DATABASE_URL in environment - load from there
if os.getenv("DATABASE_URL", None) is not None and prisma_client is None: if prisma_client is None:
prisma_setup(database_url=os.getenv("DATABASE_URL")) prisma_setup(database_url=os.getenv("DATABASE_URL"))
### LOAD CONFIG ### ### LOAD CONFIG ###
@ -1184,7 +1184,7 @@ async def startup_event():
) # start the background health check coroutine. ) # start the background health check coroutine.
verbose_proxy_logger.debug(f"prisma client - {prisma_client}") verbose_proxy_logger.debug(f"prisma client - {prisma_client}")
if prisma_client: if prisma_client is not None:
await prisma_client.connect() 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:

3
litellm/proxy/utils.py
View file

@ -555,7 +555,8 @@ class PrismaClient:
) )
async def connect(self): async def connect(self):
try: try:
await self.db.connect() if self.db.is_connected() == False:
await self.db.connect()
except Exception as e: except Exception as e:
asyncio.create_task( asyncio.create_task(
self.proxy_logging_obj.failure_handler(original_exception=e) self.proxy_logging_obj.failure_handler(original_exception=e)

5
litellm/tests/test_proxy_startup.py
View file

@ -19,7 +19,8 @@ from litellm.proxy.proxy_server import (
save_worker_config, save_worker_config,
initialize, initialize,
startup_event, startup_event,
llm_model_list llm_model_list,
shutdown_event
) )
def test_proxy_gunicorn_startup_direct_config(): def test_proxy_gunicorn_startup_direct_config():
@ -36,6 +37,7 @@ def test_proxy_gunicorn_startup_direct_config():
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
os.environ["WORKER_CONFIG"] = config_fp os.environ["WORKER_CONFIG"] = config_fp
asyncio.run(startup_event()) asyncio.run(startup_event())
asyncio.run(shutdown_event())
except Exception as e: except Exception as e:
if "Already connected to the query engine" in str(e): if "Already connected to the query engine" in str(e):
pass pass
@ -51,6 +53,7 @@ def test_proxy_gunicorn_startup_config_dict():
worker_config = {"config": config_fp} worker_config = {"config": config_fp}
os.environ["WORKER_CONFIG"] = json.dumps(worker_config) os.environ["WORKER_CONFIG"] = json.dumps(worker_config)
asyncio.run(startup_event()) asyncio.run(startup_event())
asyncio.run(shutdown_event())
except Exception as e: except Exception as e:
if "Already connected to the query engine" in str(e): if "Already connected to the query engine" in str(e):
pass pass