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

(fix) ProxyStartup - Check that prisma connection is healthy when starting an instance of LiteLLM (#6627)

Browse source 
* fix debug statements

* fix assert prisma_client.health_check is called on _setup

* asser that _setup_prisma_client is called on startup proxy

* fix prisma client health_check

* add test_bad_database_url

* add strict checks on db startup

* temp remove fix to validate if check works as expected

* add health_check back

* test_proxy_server_prisma_setup_invalid_db
This commit is contained in:
Ishaan Jaff 2024-11-06 17:36:48 -08:00 • committed by GitHub
parent 8a2b6fd8d2
commit 373f9d409e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 86 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/utils.py
View file

@ -1083,19 +1083,16 @@ class PrismaClient:
proxy_logging_obj: ProxyLogging,
http_client: Optional[Any] = None,
):
verbose_proxy_logger.debug(
"LiteLLM: DATABASE_URL Set in config, trying to 'pip install prisma'"
)
## init logging object
self.proxy_logging_obj = proxy_logging_obj
self.iam_token_db_auth: Optional[bool] = str_to_bool(
os.getenv("IAM_TOKEN_DB_AUTH")
)
verbose_proxy_logger.debug("Creating Prisma Client..")
try:
from prisma import Prisma # type: ignore
except Exception:
raise Exception("Unable to find Prisma binaries.")
verbose_proxy_logger.debug("Connecting Prisma Client to DB..")
if http_client is not None:
self.db = PrismaWrapper(
original_prisma=Prisma(http=http_client),
@ -1114,7 +1111,7 @@ class PrismaClient:
else False
),
) # Client to connect to Prisma db
verbose_proxy_logger.debug("Success - Connected Prisma Client to DB")
verbose_proxy_logger.debug("Success - Created Prisma Client")
def hash_token(self, token: str):
# Hash the string using SHA-256
@ -2348,11 +2345,7 @@ class PrismaClient:
"""
start_time = time.time()
try:
sql_query = """
SELECT 1
FROM "LiteLLM_VerificationToken"
LIMIT 1
"""
sql_query = "SELECT 1"
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments