mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(model_management): strip server-computed model_info fields on save
The model-listing endpoints inject access_via_team_ids and direct_access into each model's model_info for display. The Admin UI round-trips the whole model_info it received on save, so update_db_model and the /model/new path persisted these server-computed fields verbatim, growing model_info O(number of teams) on every UI save and leaving a stale team list that every deployment deepcopy then pays for. Strip access_via_team_ids and direct_access from model_info before persisting in both update_db_model and _add_model_to_db, cleaning up already-polluted rows on their next save.
This commit is contained in:
parent
a7d01cb1ac
commit
a9ac8b9b92
3 changed files with 108 additions and 6 deletions
|
|
@ -56,6 +56,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import
|
|||
UpdateUsefulLinksRequest,
|
||||
)
|
||||
from litellm.types.router import (
|
||||
SERVER_COMPUTED_MODEL_INFO_FIELDS,
|
||||
SPECIAL_MODEL_INFO_PARAMS,
|
||||
Deployment,
|
||||
DeploymentTypedDict,
|
||||
|
|
@ -153,10 +154,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
|
|||
prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"])
|
||||
|
||||
if "model_info" in merged_deployment_dict:
|
||||
model_info = merged_deployment_dict["model_info"]
|
||||
for key, value in model_info.items():
|
||||
if isinstance(value, datetime.datetime):
|
||||
model_info[key] = value.isoformat()
|
||||
model_info = {
|
||||
key: (value.isoformat() if isinstance(value, datetime.datetime) else value)
|
||||
for key, value in merged_deployment_dict["model_info"].items()
|
||||
if key not in SERVER_COMPUTED_MODEL_INFO_FIELDS
|
||||
}
|
||||
prisma_compatible_model_dict["model_info"] = json.dumps(model_info)
|
||||
|
||||
if updated_patch.blocked is not None:
|
||||
|
|
@ -484,8 +486,12 @@ async def _add_model_to_db(
|
|||
"model_id": model_params.model_info.id,
|
||||
"model_name": model_params.model_name,
|
||||
"litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), # type: ignore
|
||||
"model_info": model_params.model_info.model_dump_json( # type: ignore
|
||||
exclude_none=True
|
||||
"model_info": json.dumps(
|
||||
{
|
||||
k: v
|
||||
for k, v in model_params.model_info.model_dump(mode="json", exclude_none=True).items()
|
||||
if k not in SERVER_COMPUTED_MODEL_INFO_FIELDS
|
||||
}
|
||||
),
|
||||
"created_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
|
|
|
|||
|
|
@ -444,6 +444,8 @@ SPECIAL_MODEL_INFO_PARAMS = [
|
|||
"cache_creation_input_token_cost",
|
||||
]
|
||||
|
||||
SERVER_COMPUTED_MODEL_INFO_FIELDS = frozenset({"access_via_team_ids", "direct_access"})
|
||||
|
||||
|
||||
class Deployment(BaseModel):
|
||||
model_name: str
|
||||
|
|
|
|||
|
|
@ -2938,6 +2938,100 @@ class TestUpdateDBModelClearPricing:
|
|||
assert info["cache_creation_input_token_cost"] == 0.000003
|
||||
|
||||
|
||||
class TestStripServerComputedModelInfoOnSave:
|
||||
"""`access_via_team_ids` and `direct_access` are computed per-request by the
|
||||
model-listing endpoints for display. The Admin UI round-trips the whole
|
||||
`model_info` it received on save, so both the update and create paths must
|
||||
strip these server-computed fields before persisting - otherwise
|
||||
`model_info` grows O(number of teams) on every save and goes stale.
|
||||
"""
|
||||
|
||||
def test_update_db_model_strips_server_computed_fields(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
update_db_model,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
result = update_db_model(
|
||||
db_model=_build_db_model_for_blocked_test(),
|
||||
updated_patch=updateDeployment(
|
||||
model_info=ModelInfo(
|
||||
id="dep-0",
|
||||
access_via_team_ids=[f"team-{i}" for i in range(1438)],
|
||||
direct_access=True,
|
||||
base_model="azure/gpt-4o",
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
info = json.loads(result["model_info"])
|
||||
assert "access_via_team_ids" not in info
|
||||
assert "direct_access" not in info
|
||||
# legitimate, admin-supplied model_info survives
|
||||
assert info["base_model"] == "azure/gpt-4o"
|
||||
|
||||
def test_update_db_model_strips_fields_already_in_db_row(self):
|
||||
"""A row polluted by a previous UI save is cleaned up on the next save,
|
||||
even when the patch itself doesn't mention the stale fields."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
update_db_model,
|
||||
)
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params=LiteLLM_Params(model="openai/*"),
|
||||
model_info=ModelInfo(
|
||||
id="dep-stale",
|
||||
access_via_team_ids=["team-a", "team-b"],
|
||||
direct_access=False,
|
||||
),
|
||||
)
|
||||
|
||||
result = update_db_model(
|
||||
db_model=db_model,
|
||||
updated_patch=updateDeployment(model_name="openai/*"),
|
||||
)
|
||||
|
||||
info = json.loads(result["model_info"])
|
||||
assert "access_via_team_ids" not in info
|
||||
assert "direct_access" not in info
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_model_to_db_strips_server_computed_fields(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_add_model_to_db,
|
||||
)
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
model_params = Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params=LiteLLM_Params(model="openai/*"),
|
||||
model_info=ModelInfo(
|
||||
id="dep-new",
|
||||
access_via_team_ids=["team-a", "team-b", "team-c"],
|
||||
direct_access=True,
|
||||
base_model="azure/gpt-4o",
|
||||
),
|
||||
)
|
||||
|
||||
model_response = await _add_model_to_db(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
),
|
||||
prisma_client=MagicMock(),
|
||||
new_encryption_key="sk-test-salt-key",
|
||||
should_create_model_in_db=False,
|
||||
)
|
||||
|
||||
assert model_response is not None
|
||||
assert model_response.model_info is not None
|
||||
assert "access_via_team_ids" not in model_response.model_info
|
||||
assert "direct_access" not in model_response.model_info
|
||||
assert model_response.model_info["base_model"] == "azure/gpt-4o"
|
||||
|
||||
|
||||
class TestGetModelInfoWithIdBlocked:
|
||||
"""`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked`
|
||||
column into the in-memory `model_info` dict so the router filter can read it."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue