fix(proxy): stop team BYOK model name corruption on model edit (#29731)
* fix(proxy): stop team model name corruption on edit (#28382) (#29001) Team-scoped ("Team-BYOK") models store an internal routing key model_name_{team_id}_{uuid} in the model_name column and the user-facing name in model_info.team_public_model_name. The internal name leaked into /v1, /v2, and /model/info responses; the dashboard bound its edit form to it, so any non-rename save (e.g. a TPM tweak) PATCHed the internal name back. The update path then treated it as a rename, overwriting team_public_model_name and rewriting the team's models[] ACL with the mangled string -- breaking team key calls with team_model_access_denied. Two-layer fix: - Read path (root cause): add _translate_model_name_for_response and apply it in model_info_v2 and _get_proxy_model_info so /v1, /v2, and /model/info surface the public name for team-scoped rows. The DB column and router index keep the internal name as the routing key; this is a presentation-layer swap on a shallow copy (never mutates input). - Write path (defense in depth): harden _get_public_model_name so a value matching the internal shape, or a no-op against the current DB column, is never treated as a rename -- for both the top-level model_name and an explicit model_info.team_public_model_name. Tests: regression for the reported scenario, full branch coverage of _get_public_model_name, two internal-shape guard cases, an end-to-end PATCH through _update_team_model_in_db (asserts the team ACL is untouched), and four response-translation cases. 60 passed (model management), 181 passed (proxy server). * fix(ui): key Agent Builder agent selection on model_info.id (#29729) * fix(ui): key Agent Builder agent selection on model_info.id Once team-scoped BYOK models can share a public name (the backend now returns the public name on /model/info instead of the internal routing key), selecting agents by model_name collides. Key selection, create, update and delete on the stable model_info.id instead, falling back to model_name only for config-defined agents that have no id. * fix(ui): add name-match fallback to post-create agent selection If the just-created agent's id is not yet present in the re-fetched list, try matching by name before falling back to the first agent. Addresses greptile review on #29729. --------- Co-authored-by: tushar8408 <32977767+tushar8408@users.noreply.github.com>
This commit is contained in:
yuneng-jiang
GitHub
parent
f3811ce63b
commit
56aa55b991
@ -490,9 +490,45 @@ def _get_public_model_name(
|
||||
patch_data: updateDeployment,
|
||||
db_model: Deployment,
|
||||
) -> str:
|
||||
"""Determine the public model name from patch or existing model."""
|
||||
if patch_data.model_name:
|
||||
return patch_data.model_name
|
||||
"""Determine the public model name from patch or existing model.
|
||||
|
||||
The top-level ``model_name`` is the rename channel. For team-scoped rows
|
||||
the DB ``model_name`` column holds an internal routing key
|
||||
(``model_name_{team_id}_{uuid}``), and ``/model/info`` historically leaked
|
||||
it into the dashboard edit form, so a non-rename save (e.g. a TPM tweak)
|
||||
would PATCH the internal name and the update path would treat it as a
|
||||
rename -- overwriting ``team_public_model_name`` and rewriting the team ACL
|
||||
(see issue #28382).
|
||||
|
||||
Guard against that by ignoring an incoming ``model_name`` that matches the
|
||||
internal shape, or is a no-op against the current DB column. Anything else
|
||||
is a genuine rename and wins. We deliberately do NOT read
|
||||
``patch_data.model_info.team_public_model_name``: the dashboard passes the
|
||||
existing ``model_info`` blob through untouched on a rename, so honoring it
|
||||
would return the OLD public name and silently drop the rename.
|
||||
|
||||
Precedence (highest first):
|
||||
1. patch_data.model_name -- a genuine rename: not internal-shape and not a
|
||||
no-op against db_model.model_name.
|
||||
2. db_model.model_info.team_public_model_name -- existing public name.
|
||||
3. db_model.model_name -- last-resort fallback for legacy rows.
|
||||
"""
|
||||
team_id = (patch_data.model_info.team_id if patch_data.model_info else None) or (
|
||||
db_model.model_info.team_id if db_model.model_info else None
|
||||
)
|
||||
|
||||
def _is_internal_shape(name: Optional[str]) -> bool:
|
||||
if team_id is None or not name:
|
||||
return False
|
||||
return name.startswith(f"model_name_{team_id}_")
|
||||
|
||||
incoming = patch_data.model_name
|
||||
if (
|
||||
incoming
|
||||
and not _is_internal_shape(incoming)
|
||||
and incoming != db_model.model_name
|
||||
):
|
||||
return incoming
|
||||
|
||||
if db_model.model_info and db_model.model_info.team_public_model_name:
|
||||
return db_model.model_info.team_public_model_name
|
||||
|
||||
@ -11888,6 +11888,9 @@ async def model_info_v2(
|
||||
# Update total count to include agents
|
||||
search_total_count = len(all_models)
|
||||
|
||||
# Translate `model_name` to the public name for team-scoped rows.
|
||||
all_models = [_translate_model_name_for_response(m) for m in all_models]
|
||||
|
||||
return _paginate_models_response(
|
||||
all_models=all_models,
|
||||
page=page,
|
||||
@ -12322,6 +12325,33 @@ async def model_metrics_exceptions(
|
||||
return {"data": response, "exception_types": list(exception_types)}
|
||||
|
||||
|
||||
def _translate_model_name_for_response(model: dict) -> dict:
|
||||
"""For team-scoped DB rows, replace `model_name` with the public name
|
||||
in `model_info.team_public_model_name` before returning. The DB column
|
||||
and the in-memory router index keep the internal mangled name
|
||||
(`model_name_{team_id}_{uuid}`) as the routing key -- this swap is a
|
||||
presentation-layer concern. Returns a shallow copy; never mutates.
|
||||
|
||||
Without this swap the internal name leaks into `/v1/model/info` and
|
||||
`/v2/model/info`, the dashboard binds its edit form to it, and a
|
||||
non-rename save round-trips the internal name back -- corrupting
|
||||
`team_public_model_name` and the team ACL (see issue #28382).
|
||||
"""
|
||||
if not isinstance(model, dict):
|
||||
return model
|
||||
model_info = model.get("model_info") or {}
|
||||
if not isinstance(model_info, dict):
|
||||
return model
|
||||
team_public = model_info.get("team_public_model_name")
|
||||
team_id = model_info.get("team_id")
|
||||
if not team_public or not team_id:
|
||||
return model
|
||||
current = model.get("model_name") or ""
|
||||
if not current.startswith(f"model_name_{team_id}_"):
|
||||
return model
|
||||
return {**model, "model_name": team_public}
|
||||
|
||||
|
||||
def _get_proxy_model_info(model: dict) -> dict:
|
||||
# provided model_info in config.yaml
|
||||
model_info = model.get("model_info", {})
|
||||
@ -12362,7 +12392,7 @@ def _get_proxy_model_info(model: dict) -> dict:
|
||||
deployment_dict=model, excluded_keys={"litellm_credential_name"}
|
||||
)
|
||||
|
||||
return model
|
||||
return _translate_model_name_for_response(model)
|
||||
|
||||
|
||||
@router.get(
|
||||
@ -12502,8 +12532,11 @@ async def model_info_v1( # noqa: PLR0915
|
||||
else:
|
||||
all_models = []
|
||||
|
||||
for in_place_model in all_models:
|
||||
in_place_model = _get_proxy_model_info(model=in_place_model)
|
||||
# Reassign each entry: _get_proxy_model_info returns a (possibly new)
|
||||
# dict via _translate_model_name_for_response, which does NOT mutate in
|
||||
# place. Binding only the loop variable would drop the public-name swap
|
||||
# for team-scoped rows and leak the internal routing key (#28382).
|
||||
all_models = [_get_proxy_model_info(model=model) for model in all_models]
|
||||
|
||||
verbose_proxy_logger.debug("all_models: %s", all_models)
|
||||
return {"data": all_models}
|
||||
|
||||
@ -1129,6 +1129,307 @@ class TestTeamModelUpdate:
|
||||
)
|
||||
assert "403" in str(exc_info.value)
|
||||
|
||||
def test_get_public_model_name_28382_dashboard_echo_preserves_public_name(self):
|
||||
"""Regression for #28382 - a non-rename dashboard PATCH echoes the
|
||||
internal generated model_name (model_name_{team}_{uuid}) at the top
|
||||
level. That internal-shape value must be ignored (not treated as a
|
||||
rename), so _get_public_model_name falls through to the existing public
|
||||
name instead of overwriting it with the internal one."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_preserves_db_public_name_when_internal_name_unchanged(
|
||||
self,
|
||||
):
|
||||
"""If patch_data.model_info has no team_public_model_name and
|
||||
patch_data.model_name equals db_model.model_name (dashboard re-sending
|
||||
the internal name without touching the public-name field), the
|
||||
existing db_model.model_info.team_public_model_name must be preserved."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_allows_top_level_rename(self):
|
||||
"""A genuine rename via the top-level model_name field (no
|
||||
patch_data.model_info.team_public_model_name supplied, and the new
|
||||
name differs from the existing internal db model_name) must still
|
||||
return the new name."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="old-public-name",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="new-public-name",
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "new-public-name"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_top_level_rename_wins_over_stale_model_info(self):
|
||||
"""Regression (codex review): on a dashboard rename the UI sends the new
|
||||
name in model_name but passes the existing model_info blob through
|
||||
untouched -- so it still carries the OLD team_public_model_name. The
|
||||
top-level rename must win; otherwise _update_existing_team_model_assignment
|
||||
sees no change, never updates the team ACL, and the rename is silently
|
||||
dropped while the UI optimistically shows the new name."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_team-a_abc123",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-4.1"),
|
||||
model_info=ModelInfo(
|
||||
team_id="team-a", team_public_model_name="old-public-name"
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="new-public-name",
|
||||
model_info=ModelInfo(
|
||||
team_id="team-a",
|
||||
team_public_model_name="old-public-name", # stale, untouched by UI
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "new-public-name"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_falls_back_to_db_public_name(self):
|
||||
"""When patch_data carries no name hints at all (neither model_name
|
||||
nor model_info.team_public_model_name), fall back to the existing
|
||||
db_model.model_info.team_public_model_name."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_last_resort_returns_db_model_name(self):
|
||||
"""Legacy rows may have no team_public_model_name anywhere; the
|
||||
function must still return a string (the existing db_model.model_name)
|
||||
rather than raising."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="legacy-model",
|
||||
litellm_params=LiteLLM_Params(model="azure/legacy"),
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "legacy-model"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_ignores_different_internal_shape_name(self):
|
||||
"""A stale client may PATCH an internal-shaped model_name that does not
|
||||
equal the current DB column (e.g. a different uuid). It must NOT be
|
||||
treated as a rename -- fall through to the existing public name."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_realuuid",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="model_name_test-team_differentuuid",
|
||||
model_info=ModelInfo(team_id="test-team"),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
def test_get_public_model_name_ignores_internal_shape_patch_public(self):
|
||||
"""If a corrupted row round-trips an internal-shaped value in
|
||||
model_info.team_public_model_name, it must not be accepted as the
|
||||
public name -- fall through to the existing db public name."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_get_public_model_name,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_realuuid",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2-low-rpm-testing"),
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_info=ModelInfo(
|
||||
team_id="test-team",
|
||||
team_public_model_name="model_name_test-team_realuuid",
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
_get_public_model_name(patch_data=patch_data, db_model=db_model)
|
||||
== "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dashboard_edit_preserves_public_name_and_acl(self):
|
||||
"""End-to-end regression for #28382: PATCH payload shaped like the
|
||||
dashboard's model-edit form (top-level model_name = internal generated
|
||||
name, model_info.team_public_model_name = public name) must NOT trigger
|
||||
a public-name rename, must NOT touch the team ACL, and must serialize
|
||||
the public name back into model_info."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_update_team_model_in_db,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="azure/gpt-5.2-low-rpm-testing",
|
||||
custom_llm_provider="azure",
|
||||
),
|
||||
model_info=ModelInfo(
|
||||
id="model-id-123",
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="model_name_test-team_abc123",
|
||||
litellm_params=None,
|
||||
model_info=ModelInfo(
|
||||
id="model-id-123",
|
||||
team_id="test-team",
|
||||
team_public_model_name="gpt-5.2-low-rpm-testing",
|
||||
),
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="test_user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
prisma_client = MockPrismaClient(team_exists=True)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add"
|
||||
) as mock_team_model_add,
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete"
|
||||
) as mock_team_model_delete,
|
||||
):
|
||||
result = await _update_team_model_in_db(
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
)
|
||||
|
||||
# team ACL must not be touched on a no-op edit
|
||||
mock_team_model_add.assert_not_called()
|
||||
mock_team_model_delete.assert_not_called()
|
||||
|
||||
# the merged model_info written to the DB must keep the public name
|
||||
model_info_json = result.get("model_info", "")
|
||||
parsed_model_info = json.loads(model_info_json)
|
||||
assert (
|
||||
parsed_model_info.get("team_public_model_name") == "gpt-5.2-low-rpm-testing"
|
||||
)
|
||||
|
||||
# the internal model_name must not have been overwritten (caller
|
||||
# intentionally clears patch_data.model_name so the DB row's name
|
||||
# column is left alone)
|
||||
assert result.get("model_name") == "model_name_test-team_abc123"
|
||||
|
||||
|
||||
class TestModelInfoEndpoint:
|
||||
"""Test the model_info endpoint for retrieving individual model information"""
|
||||
|
||||
@ -0,0 +1,178 @@
|
||||
"""Coverage for team-scoped model-name translation in /model/info responses.
|
||||
|
||||
These live in tests/test_litellm/proxy/proxy_server/ (not the top-level
|
||||
test_proxy_server.py) because the CI coverage job collects this directory.
|
||||
They exercise the read-path fix for issue #28382: `/v1`, `/v2`, and
|
||||
`/model/info` must surface `model_info.team_public_model_name` for team-scoped
|
||||
rows instead of the internal routing key `model_name_{team_id}_{uuid}`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import (
|
||||
_get_proxy_model_info,
|
||||
_translate_model_name_for_response,
|
||||
)
|
||||
|
||||
|
||||
def _team_row() -> dict:
|
||||
return {
|
||||
"model_name": "model_name_team-abc-123_4a6b8",
|
||||
"litellm_params": {"model": "azure/gpt-5.2-low-rpm-testing"},
|
||||
"model_info": {
|
||||
"id": "byok-id-1",
|
||||
"team_id": "team-abc-123",
|
||||
"team_public_model_name": "team-claude-sonnet",
|
||||
"db_model": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_translate_swaps_internal_name_for_public():
|
||||
"""Team-scoped row: model_name is swapped to the public name."""
|
||||
result = _translate_model_name_for_response(_team_row())
|
||||
assert result["model_name"] == "team-claude-sonnet"
|
||||
|
||||
|
||||
def test_translate_leaves_global_row_untouched():
|
||||
"""No team_id / team_public_model_name -> pass through unchanged."""
|
||||
model = {
|
||||
"model_name": "gpt-4o",
|
||||
"litellm_params": {"model": "gpt-4o"},
|
||||
"model_info": {"id": "normal-id-1", "db_model": False},
|
||||
}
|
||||
assert _translate_model_name_for_response(model)["model_name"] == "gpt-4o"
|
||||
|
||||
|
||||
def test_translate_leaves_non_internal_shape_untouched():
|
||||
"""Team row whose model_name is not the internal routing key is not rewritten."""
|
||||
model = _team_row()
|
||||
model["model_name"] = "already-public-name"
|
||||
assert (
|
||||
_translate_model_name_for_response(model)["model_name"] == "already-public-name"
|
||||
)
|
||||
|
||||
|
||||
def test_translate_handles_missing_or_non_dict_model_info():
|
||||
"""Missing / None / non-dict model_info, and a non-dict model, must not raise."""
|
||||
# missing model_info
|
||||
assert _translate_model_name_for_response({"model_name": "x"})["model_name"] == "x"
|
||||
# model_info is None -> coerced to {} -> no team fields
|
||||
assert (
|
||||
_translate_model_name_for_response({"model_name": "x", "model_info": None})[
|
||||
"model_name"
|
||||
]
|
||||
== "x"
|
||||
)
|
||||
# model_info is a truthy non-dict (e.g. a stray string) -> early return
|
||||
assert (
|
||||
_translate_model_name_for_response(
|
||||
{"model_name": "x", "model_info": "garbage"}
|
||||
)["model_name"]
|
||||
== "x"
|
||||
)
|
||||
# model itself is not a dict
|
||||
assert _translate_model_name_for_response("not-a-dict") == "not-a-dict" # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_translate_does_not_mutate_input():
|
||||
"""Returns a shallow copy; the router's in-memory list keeps the routing key."""
|
||||
model = _team_row()
|
||||
result = _translate_model_name_for_response(model)
|
||||
assert result is not model
|
||||
assert model["model_name"] == "model_name_team-abc-123_4a6b8"
|
||||
|
||||
|
||||
def test_get_proxy_model_info_returns_public_name_for_team_row():
|
||||
"""`_get_proxy_model_info` must return the public name for a team-scoped
|
||||
row. Because _translate_model_name_for_response returns a shallow copy
|
||||
(it does not mutate), callers MUST use the return value -- the
|
||||
`/v1/model/info` list path historically discarded it, leaking the internal
|
||||
routing key (#28382)."""
|
||||
# Mirror the (fixed) /v1/model/info list path: assign the return back.
|
||||
all_models = [_get_proxy_model_info(model=m) for m in [_team_row()]]
|
||||
assert all_models[0]["model_name"] == "team-claude-sonnet"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v2_translates_team_model_name(monkeypatch):
|
||||
"""/v2/model/info must surface the public name for team-scoped rows.
|
||||
Covers the translation step in model_info_v2 (the read-path call site)."""
|
||||
router = MagicMock()
|
||||
router.model_list = [_team_row()]
|
||||
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(ps.proxy_config, "get_config", AsyncMock(return_value={}))
|
||||
monkeypatch.setattr(
|
||||
ps,
|
||||
"_apply_search_filter_to_models",
|
||||
AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
import litellm.proxy.agent_endpoints.model_list_helpers as mlh
|
||||
|
||||
monkeypatch.setattr(
|
||||
mlh,
|
||||
"append_agents_to_model_info",
|
||||
AsyncMock(side_effect=lambda models, **kw: models),
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
# Pass every query param explicitly: called directly (not through FastAPI),
|
||||
# the fastapi.Query(...) defaults are Query objects, not their values.
|
||||
resp = await ps.model_info_v2(
|
||||
user_api_key_dict=admin,
|
||||
model=None,
|
||||
user_models_only=False,
|
||||
include_team_models=False,
|
||||
debug=False,
|
||||
page=1,
|
||||
size=50,
|
||||
search=None,
|
||||
modelId=None,
|
||||
teamId=None,
|
||||
sortBy=None,
|
||||
sortOrder="asc",
|
||||
)
|
||||
|
||||
names = [m["model_name"] for m in resp["data"]]
|
||||
assert "team-claude-sonnet" in names
|
||||
assert "model_name_team-abc-123_4a6b8" not in names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch):
|
||||
"""/v1/model/info list path (no litellm_model_id) must surface the public
|
||||
name. Covers the list comprehension that assigns _get_proxy_model_info's
|
||||
return back into all_models (#28382 review)."""
|
||||
router = MagicMock()
|
||||
router.get_model_names.return_value = ["team-claude-sonnet"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
router.get_model_list.return_value = [_team_row()]
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", [_team_row()])
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "get_key_models", lambda **kw: [])
|
||||
monkeypatch.setattr(ps, "get_team_models", lambda **kw: [])
|
||||
monkeypatch.setattr(
|
||||
ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"]
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None)
|
||||
|
||||
names = [m["model_name"] for m in resp["data"]]
|
||||
assert "team-claude-sonnet" in names
|
||||
assert "model_name_team-abc-123_4a6b8" not in names
|
||||
@ -131,6 +131,14 @@ function getAgentModelId(agent: AgentModel): string | null {
|
||||
return info?.id ?? null;
|
||||
}
|
||||
|
||||
// Selection key that always resolves to a non-null string. Prefers the DB
|
||||
// id (stable across renames and unique across teams) but falls back to
|
||||
// `model_name` so config-file-defined agents — which have no `model_info.id`
|
||||
// — remain selectable.
|
||||
function getAgentSelectionKey(agent: AgentModel): string {
|
||||
return getAgentModelId(agent) ?? agent.model_name;
|
||||
}
|
||||
|
||||
function parseUnderlyingModel(litellmModel: string | undefined): string | undefined {
|
||||
if (!litellmModel || !litellmModel.startsWith("litellm_agent/")) return undefined;
|
||||
return litellmModel.slice("litellm_agent/".length) || undefined;
|
||||
@ -196,22 +204,24 @@ export default function AgentBuilderView({
|
||||
|
||||
const effectiveApiKey = apiKey || accessToken || "";
|
||||
const selectedAgent =
|
||||
selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => a.model_name === selectedId) ?? null;
|
||||
selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => getAgentSelectionKey(a) === selectedId) ?? null;
|
||||
const isNewAgent = selectedId === NEW_AGENT_ID;
|
||||
const selectedAgentModelId = selectedAgent ? getAgentModelId(selectedAgent) : null;
|
||||
|
||||
const loadAgents = useCallback(async () => {
|
||||
if (!accessToken || !userID || !userRole) return;
|
||||
const loadAgents = useCallback(async (): Promise<AgentModel[]> => {
|
||||
if (!accessToken || !userID || !userRole) return [];
|
||||
setLoadingAgents(true);
|
||||
try {
|
||||
const list = await fetchAvailableAgentModels(accessToken, userID, userRole);
|
||||
setAgentModels(list);
|
||||
if (!selectedId || (selectedId !== NEW_AGENT_ID && !list.some((a) => a.model_name === selectedId))) {
|
||||
setSelectedId(list.length > 0 ? list[0].model_name : null);
|
||||
if (!selectedId || (selectedId !== NEW_AGENT_ID && !list.some((a) => getAgentSelectionKey(a) === selectedId))) {
|
||||
setSelectedId(list.length > 0 ? getAgentSelectionKey(list[0]) : null);
|
||||
}
|
||||
return list;
|
||||
} catch (e) {
|
||||
console.error(e);
|
||||
NotificationsManager.fromBackend("Failed to load agents");
|
||||
return [];
|
||||
} finally {
|
||||
setLoadingAgents(false);
|
||||
}
|
||||
@ -308,7 +318,7 @@ export default function AgentBuilderView({
|
||||
}
|
||||
setSaving(true);
|
||||
try {
|
||||
await modelCreateCall(accessToken, {
|
||||
const response = await modelCreateCall(accessToken, {
|
||||
model_name: draftName.trim(),
|
||||
litellm_params: {
|
||||
model: `litellm_agent/${draftUnderlyingModel}`,
|
||||
@ -319,9 +329,15 @@ export default function AgentBuilderView({
|
||||
},
|
||||
model_info: {},
|
||||
});
|
||||
const newName = draftName.trim();
|
||||
await loadAgents();
|
||||
setSelectedId(newName);
|
||||
// /model/new returns the row with `model_id` at the top level.
|
||||
// Prefer that id over name-matching so we land on the just-created
|
||||
// agent even when its public name collides with another team's.
|
||||
const createdId: string | null = response?.model_id ?? response?.model_info?.id ?? null;
|
||||
const list = await loadAgents();
|
||||
const created = createdId
|
||||
? list.find((a) => getAgentModelId(a) === createdId) ?? list.find((a) => a.model_name === draftName.trim())
|
||||
: list.find((a) => a.model_name === draftName.trim());
|
||||
setSelectedId(created ? getAgentSelectionKey(created) : list[0] ? getAgentSelectionKey(list[0]) : null);
|
||||
setActiveTab("chat");
|
||||
} catch (e) {
|
||||
NotificationsManager.fromBackend("Failed to save agent");
|
||||
@ -353,8 +369,10 @@ export default function AgentBuilderView({
|
||||
selectedAgentModelId,
|
||||
);
|
||||
NotificationsManager.success("Agent updated successfully");
|
||||
await loadAgents();
|
||||
setSelectedId(draftName.trim());
|
||||
const list = await loadAgents();
|
||||
const stillSelected = list.find((a) => getAgentModelId(a) === selectedAgentModelId);
|
||||
const target = stillSelected ?? list[0];
|
||||
setSelectedId(target ? getAgentSelectionKey(target) : null);
|
||||
} catch (e) {
|
||||
NotificationsManager.fromBackend("Failed to update agent");
|
||||
} finally {
|
||||
@ -398,9 +416,9 @@ export default function AgentBuilderView({
|
||||
try {
|
||||
await modelDeleteCall(accessToken, selectedAgentModelId);
|
||||
NotificationsManager.success("Agent deleted");
|
||||
await loadAgents();
|
||||
const remaining = agentModels.filter((a) => a.model_name !== selectedAgent.model_name);
|
||||
setSelectedId(remaining.length > 0 ? remaining[0].model_name : null);
|
||||
const list = await loadAgents();
|
||||
const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId);
|
||||
setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null);
|
||||
} catch (e) {
|
||||
NotificationsManager.fromBackend("Failed to delete agent");
|
||||
} finally {
|
||||
@ -462,21 +480,24 @@ export default function AgentBuilderView({
|
||||
</div>
|
||||
) : (
|
||||
<>
|
||||
{agentModels.map((agent) => (
|
||||
<button
|
||||
key={agent.model_name}
|
||||
type="button"
|
||||
onClick={() => setSelectedId(agent.model_name)}
|
||||
className={`mb-1 w-full rounded-md border-l-2 px-3 py-2 text-left text-sm transition-colors ${
|
||||
selectedId === agent.model_name
|
||||
? "border-blue-500 bg-blue-50 text-blue-800"
|
||||
: "border-transparent hover:bg-gray-50"
|
||||
}`}
|
||||
>
|
||||
<div className="font-medium truncate">{agent.model_name}</div>
|
||||
<div className="text-[10px] text-gray-500 truncate">litellm_agent</div>
|
||||
</button>
|
||||
))}
|
||||
{agentModels.map((agent) => {
|
||||
const key = getAgentSelectionKey(agent);
|
||||
return (
|
||||
<button
|
||||
key={key}
|
||||
type="button"
|
||||
onClick={() => setSelectedId(key)}
|
||||
className={`mb-1 w-full rounded-md border-l-2 px-3 py-2 text-left text-sm transition-colors ${
|
||||
selectedId === key
|
||||
? "border-blue-500 bg-blue-50 text-blue-800"
|
||||
: "border-transparent hover:bg-gray-50"
|
||||
}`}
|
||||
>
|
||||
<div className="font-medium truncate">{agent.model_name}</div>
|
||||
<div className="text-[10px] text-gray-500 truncate">litellm_agent</div>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleAddAgent}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user