From dcd8f7ebf2fffaef060cfabf4b31c499b73c6e44 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 22 Jul 2024 18:29:21 -0700 Subject: [PATCH] control team callbacks using API --- litellm/proxy/litellm_pre_call_utils.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1014a325a..642c12616 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional from fastapi import Request from litellm._logging import verbose_logger, verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, TeamCallbackMetadata, UserAPIKeyAuth from litellm.types.utils import SupportedCacheControls if TYPE_CHECKING: @@ -207,6 +207,29 @@ async def add_litellm_data_to_request( **data, } # add the team-specific configs to the completion call + # Team Callbacks controls + if user_api_key_dict.team_metadata is not None: + team_metadata = user_api_key_dict.team_metadata + if "callback_settings" in team_metadata: + callback_settings = team_metadata.get("callback_settings", None) or {} + callback_settings_obj = TeamCallbackMetadata(**callback_settings) + """ + callback_settings = { + { + 'callback_vars': {'langfuse_public_key': 'pk', 'langfuse_secret_key': 'sk_'}, + 'failure_callback': [], + 'success_callback': ['langfuse', 'langfuse'] + } + } + """ + data["success_callback"] = callback_settings_obj.success_callback + data["failure_callback"] = callback_settings_obj.failure_callback + + if callback_settings_obj.callback_vars is not None: + # unpack callback_vars in data + for k, v in callback_settings_obj.callback_vars.items(): + data[k] = v + return data