fix - flush langfuse logs on proxy shutdown

This commit is contained in:
Ishaan Jaff 2024-05-16 09:21:48 -07:00
parent fda3914f4c
commit 579bc09c37

View file

@ -10623,6 +10623,7 @@ def _has_user_setup_sso():
@router.on_event("shutdown") @router.on_event("shutdown")
async def shutdown_event(): async def shutdown_event():
global prisma_client, master_key, user_custom_auth, user_custom_key_generate global prisma_client, master_key, user_custom_auth, user_custom_key_generate
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
if prisma_client: if prisma_client:
verbose_proxy_logger.debug("Disconnecting from Prisma") verbose_proxy_logger.debug("Disconnecting from Prisma")
await prisma_client.disconnect() await prisma_client.disconnect()
@ -10634,6 +10635,18 @@ async def shutdown_event():
if db_writer_client is not None: if db_writer_client is not None:
await db_writer_client.close() await db_writer_client.close()
# flush remaining langfuse logs
if "langfuse" in litellm.success_callback:
try:
# flush langfuse logs on shutdow
from litellm.utils import langFuseLogger
langFuseLogger.Langfuse.flush()
except:
# [DO NOT BLOCK shutdown events for this]
pass
## RESET CUSTOM VARIABLES ## ## RESET CUSTOM VARIABLES ##
cleanup_router_config_variables() cleanup_router_config_variables()