use 1 helper to merge tags + ensure unique ness

This commit is contained in:
Ishaan Jaff 2024-11-26 15:28:57 -08:00
parent 8fd3bf34d8
commit 5cdabbeb3d

View file

@ -288,12 +288,12 @@ class LiteLLMProxyRequestSetup:
## KEY-LEVEL SPEND LOGS / TAGS ## KEY-LEVEL SPEND LOGS / TAGS
if "tags" in key_metadata and key_metadata["tags"] is not None: 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"] = (
data[_metadata_variable_name]["tags"], list LiteLLMProxyRequestSetup._merge_tags(
): request_tags=data[_metadata_variable_name].get("tags"),
data[_metadata_variable_name]["tags"].extend(key_metadata["tags"]) tags_to_add=key_metadata["tags"],
else: )
data[_metadata_variable_name]["tags"] = key_metadata["tags"] )
if "spend_logs_metadata" in key_metadata and isinstance( if "spend_logs_metadata" in key_metadata and isinstance(
key_metadata["spend_logs_metadata"], dict key_metadata["spend_logs_metadata"], dict
): ):
@ -319,6 +319,28 @@ class LiteLLMProxyRequestSetup:
data["disable_fallbacks"] = key_metadata["disable_fallbacks"] data["disable_fallbacks"] = key_metadata["disable_fallbacks"]
return data return data
@staticmethod
def _merge_tags(request_tags: Optional[list], tags_to_add: Optional[list]) -> list:
"""
Helper function to merge two lists of tags, ensuring no duplicates.
Args:
request_tags (Optional[list]): List of tags from the original request
tags_to_add (Optional[list]): List of tags to add
Returns:
list: Combined list of unique tags
"""
final_tags = []
if request_tags and isinstance(request_tags, list):
final_tags.extend(request_tags)
if tags_to_add and isinstance(tags_to_add, list):
final_tags.extend(tags_to_add)
return list(set(final_tags))
async def add_litellm_data_to_request( # noqa: PLR0915 async def add_litellm_data_to_request( # noqa: PLR0915
data: dict, data: dict,
@ -442,12 +464,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915
## TEAM-LEVEL SPEND LOGS/TAGS ## TEAM-LEVEL SPEND LOGS/TAGS
team_metadata = user_api_key_dict.team_metadata or {} team_metadata = user_api_key_dict.team_metadata or {}
if "tags" in team_metadata and team_metadata["tags"] is not None: 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"] = LiteLLMProxyRequestSetup._merge_tags(
data[_metadata_variable_name]["tags"], list request_tags=data[_metadata_variable_name].get("tags"),
): tags_to_add=team_metadata["tags"],
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( if "spend_logs_metadata" in team_metadata and isinstance(
team_metadata["spend_logs_metadata"], dict team_metadata["spend_logs_metadata"], dict
): ):