From bc812672a18bdbeee3ce5e68204d2ea7a2f437a6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 7 Jan 2026 10:55:25 -0800 Subject: [PATCH] Addressing comments, pending tests --- litellm/proxy/_types.py | 2 +- .../key_management_endpoints.py | 13 +++++++------ .../proxy/management_endpoints/team_endpoints.py | 6 +++--- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2b37128713..f87ab5b4ce 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -863,7 +863,7 @@ class KeyRequestBase(GenerateRequestBase): tpm_limit_type: Optional[ Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm - router_settings: Optional[dict] = None + router_settings: Optional[UpdateRouterConfig] = None class LiteLLMKeyType(str, enum.Enum): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 387c049839..9f2bff8161 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -14,9 +14,10 @@ import copy import json import secrets import traceback +import yaml from datetime import datetime, timedelta, timezone from typing import List, Literal, Optional, Tuple, cast - +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status @@ -1390,7 +1391,7 @@ async def prepare_key_update_data( # Serialize router_settings to JSON if present if "router_settings" in non_default_values and non_default_values["router_settings"] is not None: - non_default_values["router_settings"] = json.dumps(non_default_values["router_settings"]) + non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"]) non_default_values = prepare_metadata_fields( data=data, non_default_values=non_default_values, existing_metadata=_metadata @@ -2117,7 +2118,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) - router_settings_json = json.dumps(router_settings) if router_settings is not None else json.dumps({}) + router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) # Add model_rpm_limit and model_tpm_limit to metadata if model_rpm_limit is not None: @@ -2281,9 +2282,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 router_settings_value = key_data.get("router_settings") if router_settings_value is not None and isinstance(router_settings_value, str): try: - key_data["router_settings"] = json.loads(router_settings_value) - except json.JSONDecodeError: - # If it's not valid JSON, keep as is or set to empty dict + key_data["router_settings"] = yaml.safe_load(router_settings_value) + except yaml.YAMLError: + # If it's not valid JSON/YAML, keep as is or set to empty dict key_data["router_settings"] = {} except Exception as e: verbose_proxy_logger.error( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59d68b6978..27809f57d1 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -100,7 +100,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamMemberAddResult, UpdateTeamMemberPermissionsRequest, ) - +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps router = APIRouter() @@ -904,7 +904,7 @@ async def new_team( # noqa: PLR0915 # Serialize router_settings to JSON (matching key creation pattern) router_settings_value = getattr(data, "router_settings", None) - router_settings_json = json.dumps(router_settings_value) if router_settings_value is not None else json.dumps({}) + router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) complete_team_data_dict["router_settings"] = router_settings_json complete_team_data_dict = prisma_client.jsonify_team_object( @@ -1391,7 +1391,7 @@ async def update_team( # noqa: PLR0915 # Serialize router_settings to JSON if present (matching key update pattern) if "router_settings" in updated_kv and updated_kv["router_settings"] is not None: - updated_kv["router_settings"] = json.dumps(updated_kv["router_settings"]) + updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) team_row: Optional[LiteLLM_TeamTable] = (