fix(proxy_server.py): stronger type checking for team based settings

This commit is contained in:
Krrish Dholakia 2024-02-02 13:14:23 -08:00
parent 730c5ef9f6
commit ada62192ad
2 changed files with 25 additions and 3 deletions

View file

@ -214,6 +214,13 @@ class KeyManagementSystem(enum.Enum):
LOCAL = "local" 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): class DynamoDBArgs(LiteLLMBase):
billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"] billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"]
read_capacity_units: Optional[int] = None read_capacity_units: Optional[int] = None

View file

@ -1032,6 +1032,7 @@ class ProxyConfig:
if all_teams_config is None: if all_teams_config is None:
return team_config return team_config
for team in all_teams_config: for team in all_teams_config:
if "team_id" in team:
if team_id == team["team_id"]: if team_id == team["team_id"]:
team_config = team team_config = team
break break
@ -1175,6 +1176,20 @@ class ProxyConfig:
# this is set in the cache branch # this is set in the cache branch
# see usage here: https://docs.litellm.ai/docs/proxy/caching # see usage here: https://docs.litellm.ai/docs/proxy/caching
pass 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: else:
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"