diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 9a1cf701142..bc67b788453 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -17,7 +17,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, TypedDict, overload from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view @@ -60,6 +60,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _DECODED_JSON: Final = TypeAdapter(object) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(dict[str, str]) class _TagRecord(Protocol): @@ -221,6 +222,18 @@ async def _get_model_names(prisma_client: "PrismaClient", model_ids: Sequence[st return {} +def _parse_model_info(raw_model_info: object) -> dict[str, str]: # mutable-ok: TagConfig model_info requires dict + """Helper function to parse stored model_info into a dict[str, str].""" + if not raw_model_info: + return {} + try: + if isinstance(raw_model_info, str): + return _MODEL_INFO_ADAPTER.validate_json(raw_model_info) + return _MODEL_INFO_ADAPTER.validate_python(raw_model_info) + except (ValidationError, ValueError): + return {} + + async def get_deployments_by_model(model: str, llm_router: "Router") -> list["Deployment"]: """ Get all deployments by model @@ -348,6 +361,8 @@ async def new_tag( "message": f"Tag {tag.name} created successfully", "tag": tag_config, } + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error creating tag: %s", e) raise HTTPException(status_code=500, detail=str(e)) @@ -444,15 +459,20 @@ async def update_tag( cleared_budget_fields=frozenset(field for field in tag.model_fields_set if getattr(tag, field) is None), ) - # Get model names for model_info - model_info: Final = await _get_model_names(prisma_client, tag.models or []) - # Prepare update data - update_data: Final = { - "description": tag.description, - "models": tag.models or [], - "model_info": json.dumps(model_info), - } + update_data: Final[dict[str, object]] = {} # mutable-ok: prisma update input payload built conditionally + + if "description" in tag.model_fields_set: + update_data["description"] = tag.description + + model_info: Final = ( + await _get_model_names(prisma_client, tag.models or []) + if "models" in tag.model_fields_set + else _parse_model_info(getattr(existing_tag, "model_info", None)) + ) + if "models" in tag.model_fields_set: + update_data["models"] = tag.models or [] + update_data["model_info"] = json.dumps(model_info) # Add budget_id if it changed if budget_id != existing_tag.budget_id: @@ -481,6 +501,8 @@ async def update_tag( "message": f"Tag {tag.name} updated successfully", "tag": tag_config, } + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error updating tag: %s", e) raise HTTPException(status_code=500, detail=str(e)) @@ -547,6 +569,8 @@ async def info_tag( requested_tags[tag_record.tag_name] = tag_dict return requested_tags + except HTTPException: + raise except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @@ -718,6 +742,8 @@ async def delete_tag( await _evict_tag_cache_keys((tag_cache_key(data.name), tag_registry_cache_key())) return {"message": f"Tag {data.name} deleted successfully"} + except HTTPException: + raise except Exception as e: raise HTTPException(status_code=500, detail=str(e)) diff --git a/tests/integration/spend/test_tag_budget_enforcement.py b/tests/integration/spend/test_tag_budget_enforcement.py index 0c0db2943c6..bf25cfec5c3 100644 --- a/tests/integration/spend/test_tag_budget_enforcement.py +++ b/tests/integration/spend/test_tag_budget_enforcement.py @@ -492,8 +492,6 @@ def test_tag_delete_evicts_blocked_limit( def test_tag_info_returns_404_for_deleted_tag(gateway: Gateway) -> None: - pytest.skip("BUG: POST /tag/info for a deleted tag returns 500 with detail '404: Tags not found' instead of 404") - tag: Final = f"tag-info-deleted-{uuid.uuid4().hex}" with gateway.scenario() as scenario: @@ -511,8 +509,6 @@ def test_tag_info_returns_404_for_deleted_tag(gateway: Gateway) -> None: def test_tag_partial_update_keeps_description_and_models(gateway: Gateway) -> None: - pytest.skip("BUG: /tag/update without description/models wipes the stored description and models") - tag: Final = f"tag-partial-update-{uuid.uuid4().hex}" with gateway.scenario() as scenario: model: Final = scenario.model() diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 1dfababae55..7aeff101298 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -1586,3 +1586,211 @@ def test_tag_info_and_tag_list_return_the_stored_model_info_decoded( assert info_response.json()["routed-tag"]["model_info"] == returned_model_info assert list_response.status_code == 200 assert [tag["model_info"] for tag in list_response.json()] == [returned_model_info] + + +def test_tag_info_returns_404_when_tag_not_found(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_db = Mock() + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db), + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/info", json={"names": ["nonexistent-tag"]}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 404 + assert response.json()["detail"] == "Tags not found: ['nonexistent-tag']" + + +def test_tag_update_returns_404_when_tag_not_found(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_db = Mock() + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db), + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/update", json={"name": "nonexistent-tag", "description": "new"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 404 + assert "not found" in response.json()["detail"] + + +def test_tag_delete_returns_404_when_tag_not_found(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_db = Mock() + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db), + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/delete", json={"name": "nonexistent-tag"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 404 + assert "not found" in response.json()["detail"] + + +def test_tag_partial_update_does_not_wipe_description_or_models(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + existing_tag = SimpleNamespace( + tag_name="test-tag", + description="existing description", + models=["model-1"], + model_info='{"model-1": "gpt-4"}', + budget_id=None, + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 2), + created_by="user-123", + ) + updated_tag = SimpleNamespace( + tag_name="test-tag", + description="existing description", + models=["model-1"], + model_info='{"model-1": "gpt-4"}', + budget_id="budget-1", + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 3), + created_by="user-123", + ) + + mock_db = Mock() + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + captured_update_kwargs: dict[str, object] = {} + + async def fake_update(**kwargs): + captured_update_kwargs.update(kwargs) + return updated_tag + + mock_db.litellm_tagtable.update = fake_update + + budget_record = SimpleNamespace(budget_id="budget-1", max_budget=10.0) + mock_db.litellm_budgettable.create = AsyncMock(return_value=budget_record) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db, jsonify_object=lambda data: dict(data)), + ) + monkeypatch.setattr(proxy_server, "litellm_proxy_admin_name", "admin") + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/update", json={"name": "test-tag", "max_budget": 10.0}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + res_data = response.json() + assert res_data["tag"]["description"] == "existing description" + assert res_data["tag"]["models"] == ["model-1"] + assert res_data["tag"]["model_info"] == {"model-1": "gpt-4"} + data_sent = captured_update_kwargs.get("data", {}) + assert "description" not in data_sent + assert "models" not in data_sent + assert "model_info" not in data_sent + + +def test_new_tag_duplicate_returns_400(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_db = Mock() + existing_tag = SimpleNamespace(tag_name="existing-tag") + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db), + ) + monkeypatch.setattr(proxy_server, "llm_router", Mock()) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/new", json={"name": "existing-tag"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 400 + assert "already exists" in response.json()["detail"] + + +def test_tag_update_with_invalid_model_info_returns_empty_dict(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + existing_tag = SimpleNamespace( + tag_name="test-tag", + description="existing", + models=["model-1"], + model_info="invalid-json-data", + budget_id=None, + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 2), + created_by="user-123", + ) + updated_tag = SimpleNamespace( + tag_name="test-tag", + description="existing", + models=["model-1"], + model_info="invalid-json-data", + budget_id=None, + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 3), + created_by="user-123", + ) + + mock_db = Mock() + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag) + mock_db.litellm_tagtable.update = AsyncMock(return_value=updated_tag) + + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=mock_db, jsonify_object=lambda data: dict(data)), + ) + monkeypatch.setattr(proxy_server, "litellm_proxy_admin_name", "admin") + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.post("/tag/update", json={"name": "test-tag", "description": "new-desc"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert response.json()["tag"]["model_info"] == {}