From ffeb5ce22a306cd3747d51a80836dfcfde66a8e7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 29 Aug 2024 13:05:00 -0700 Subject: [PATCH] add set / update tags for a team --- .../proxy/management_endpoints/team_endpoints.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 5b6885ecd..0858e6089 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -224,6 +224,13 @@ async def new_team( model_id=_model_id, ) + # Set tags on the new team + if data.tags is not None: + if complete_team_data.metadata is None: + complete_team_data.metadata = {"tags": data.tags} + else: + complete_team_data.metadata["tags"] = data.tags + # If budget_duration is set, set `budget_reset_at` if complete_team_data.budget_duration is not None: duration_s = _duration_in_seconds(duration=complete_team_data.budget_duration) @@ -365,6 +372,15 @@ async def update_team( # set the budget_reset_at in DB updated_kv["budget_reset_at"] = reset_at + # check if user is trying to update tags for team + if "tags" in updated_kv and updated_kv["tags"] is not None: + # remove tags from updated_kv + _tags = updated_kv.pop("tags") + if "metadata" in updated_kv and updated_kv["metadata"] is not None: + updated_kv["metadata"]["tags"] = _tags + else: + updated_kv["metadata"] = {"tags": _tags} + updated_kv = prisma_client.jsonify_object(data=updated_kv) team_row: Optional[ LiteLLM_TeamTable