feat(litellm_pre_call_utils.py): support passing tags/spend logs metadata from keys/team metadata to request

This commit is contained in:
Krrish Dholakia 2024-08-21 08:12:45 -07:00
parent 09aacfc6a8
commit 04fc0bd7b3
2 changed files with 146 additions and 1 deletions

View file

@ -224,7 +224,7 @@ async def add_litellm_data_to_request(
user_api_key_dict, "team_alias", None
)
### KEY-LEVEL Contorls
### KEY-LEVEL Controls
key_metadata = user_api_key_dict.metadata
if "cache" in key_metadata:
data["cache"] = {}
@ -233,6 +233,51 @@ async def add_litellm_data_to_request(
if k in SupportedCacheControls:
data["cache"][k] = v
## KEY-LEVEL SPEND LOGS / TAGS
if "tags" in key_metadata and key_metadata["tags"] is not None:
if "tags" in data[_metadata_variable_name] and isinstance(
data[_metadata_variable_name]["tags"], list
):
data[_metadata_variable_name]["tags"].extend(key_metadata["tags"])
else:
data[_metadata_variable_name]["tags"] = key_metadata["tags"]
if "spend_logs_metadata" in key_metadata and isinstance(
key_metadata["spend_logs_metadata"], dict
):
if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance(
data[_metadata_variable_name]["spend_logs_metadata"], dict
):
data[_metadata_variable_name]["spend_logs_metadata"].update(
key_metadata["spend_logs_metadata"]
)
else:
data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata[
"spend_logs_metadata"
]
## TEAM-LEVEL SPEND LOGS/TAGS
team_metadata = user_api_key_dict.team_metadata or {}
if "tags" in team_metadata and team_metadata["tags"] is not None:
if "tags" in data[_metadata_variable_name] and isinstance(
data[_metadata_variable_name]["tags"], list
):
data[_metadata_variable_name]["tags"].extend(team_metadata["tags"])
else:
data[_metadata_variable_name]["tags"] = team_metadata["tags"]
if "spend_logs_metadata" in team_metadata and isinstance(
team_metadata["spend_logs_metadata"], dict
):
if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance(
data[_metadata_variable_name]["spend_logs_metadata"], dict
):
data[_metadata_variable_name]["spend_logs_metadata"].update(
team_metadata["spend_logs_metadata"]
)
else:
data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata[
"spend_logs_metadata"
]
# Team spend, budget - used by prometheus.py
data[_metadata_variable_name][
"user_api_key_team_max_budget"