From ada62192ad0c874ac57ed4f320af99a415fe72d9 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 2 Feb 2024 13:14:23 -0800
Subject: [PATCH] fix(proxy_server.py): stronger type checking for team based
 settings

---
 litellm/proxy/_types.py       |  7 +++++++
 litellm/proxy/proxy_server.py | 21 ++++++++++++++++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index eb70cc07c..f7d1e8ad0 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -214,6 +214,13 @@ class KeyManagementSystem(enum.Enum):
     LOCAL = "local"
 
 
+class TeamDefaultSettings(LiteLLMBase):
+    team_id: str
+
+    class Config:
+        extra = "allow"  # allow params not defined here, these fall in litellm.completion(**kwargs)
+
+
 class DynamoDBArgs(LiteLLMBase):
     billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"]
     read_capacity_units: Optional[int] = None
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index fed61cf89..9fe25ed61 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1032,9 +1032,10 @@ class ProxyConfig:
         if all_teams_config is None:
             return team_config
         for team in all_teams_config:
-            if team_id == team["team_id"]:
-                team_config = team
-                break
+            if "team_id" in team:
+                if team_id == team["team_id"]:
+                    team_config = team
+                    break
         for k, v in team_config.items():
             if isinstance(v, str) and v.startswith("os.environ/"):
                 team_config[k] = litellm.get_secret(v)
@@ -1175,6 +1176,20 @@ class ProxyConfig:
                     # this is set in the cache branch
                     # see usage here: https://docs.litellm.ai/docs/proxy/caching
                     pass
+                elif key == "default_team_settings":
+                    for idx, team_setting in enumerate(
+                        value
+                    ):  # run through pydantic validation
+                        try:
+                            TeamDefaultSettings(**team_setting)
+                        except:
+                            raise Exception(
+                                f"team_id missing from default_team_settings at index={idx}\npassed in value={team_setting}"
+                            )
+                    verbose_proxy_logger.debug(
+                        f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"
+                    )
+                    setattr(litellm, key, value)
                 else:
                     verbose_proxy_logger.debug(
                         f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"