diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 95b7300992c..cde9f0a78bb 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -26,7 +26,6 @@ ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.types.tag_management import ( - LiteLLM_DailyTagSpendTable, TagConfig, TagDeleteRequest, TagInfoRequest, @@ -445,25 +444,30 @@ async def list_tags( list_of_tags.append(tag_dict) ## QUERY DYNAMIC TAGS ## - dynamic_tags = await prisma_client.db.litellm_dailytagspend.find_many( - distinct=["tag"], + # Use group_by instead of find_many(distinct=["tag"]). + # Prisma's distinct fetches all columns for all rows and deduplicates + # in application code, which is extremely slow on large tables. + # See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood + dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( + by=["tag"], + where={"tag": {"not": None}}, + # The old find_many(distinct=...) returned arbitrary timestamps from + # whichever row Prisma happened to pick. MIN/MAX give more meaningful + # values: earliest appearance and most recent activity. + _min={"created_at": True}, + _max={"updated_at": True}, ) - dynamic_tags_list = [ - LiteLLM_DailyTagSpendTable(**dynamic_tag.model_dump()) - for dynamic_tag in dynamic_tags - ] - dynamic_tag_config = [ { - "name": tag.tag, + "name": row["tag"], "description": "This is just a spend tag that was passed dynamically in a request. It does not control any LLM models.", "models": None, - "created_at": tag.created_at.isoformat(), - "updated_at": tag.updated_at.isoformat(), + "created_at": row["_min"]["created_at"].isoformat(), + "updated_at": row["_max"]["updated_at"].isoformat(), } - for tag in dynamic_tags_list - if tag.tag not in stored_tag_names + for row in dynamic_tag_rows + if row["tag"] not in stored_tag_names ] return list_of_tags + dynamic_tag_config diff --git a/litellm/types/tag_management.py b/litellm/types/tag_management.py index a9b58ace02f..9c1509e8167 100644 --- a/litellm/types/tag_management.py +++ b/litellm/types/tag_management.py @@ -1,4 +1,3 @@ -from datetime import datetime from typing import Dict, List, Optional from pydantic import BaseModel @@ -49,21 +48,3 @@ class TagInfoRequest(BaseModel): names: List[str] -class LiteLLM_DailyTagSpendTable(BaseModel): - id: str - tag: str - date: str - api_key: str - model: str - model_group: Optional[str] - custom_llm_provider: Optional[str] - prompt_tokens: int - completion_tokens: int - cache_read_input_tokens: int - cache_creation_input_tokens: int - spend: float - api_requests: int - successful_requests: int - failed_requests: int - created_at: datetime - updated_at: datetime diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 6da3d1f918d..f330b40282c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -244,6 +244,120 @@ async def test_delete_tag(): app.dependency_overrides.clear() +@pytest.mark.asyncio +async def test_list_tags_with_dynamic_tags(): + """ + Test that list_tags uses group_by to get distinct dynamic tags efficiently + and merges them with stored tags, excluding duplicates. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + + # Setup stored tags + stored_tag = Mock() + stored_tag.tag_name = "stored-tag" + stored_tag.description = "A stored tag" + stored_tag.models = ["model-1"] + stored_tag.model_info = {} + stored_tag.spend = 0.0 + stored_tag.budget_id = None + stored_tag.created_at = datetime(2025, 1, 1) + stored_tag.updated_at = datetime(2025, 1, 1) + stored_tag.created_by = "user-123" + stored_tag.litellm_budget_table = None + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[stored_tag]) + + # Setup dynamic tags via group_by — includes one that overlaps with stored + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[ + {"tag": "dynamic-tag-1", "_min": {"created_at": datetime(2025, 2, 1)}, "_max": {"updated_at": datetime(2025, 3, 1)}}, + {"tag": "dynamic-tag-2", "_min": {"created_at": datetime(2025, 2, 2)}, "_max": {"updated_at": datetime(2025, 3, 2)}}, + {"tag": "stored-tag", "_min": {"created_at": datetime(2025, 1, 1)}, "_max": {"updated_at": datetime(2025, 1, 1)}}, # duplicate, should be excluded + ]) + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get("/tag/list", headers=headers) + + assert response.status_code == 200 + result = response.json() + + # Should have 1 stored + 2 dynamic (the duplicate excluded) + assert len(result) == 3 + + tag_names = [t["name"] for t in result] + assert "stored-tag" in tag_names + assert "dynamic-tag-1" in tag_names + assert "dynamic-tag-2" in tag_names + + # Verify dynamic tags include created_at/updated_at + dynamic_tags = {t["name"]: t for t in result if t["name"].startswith("dynamic-")} + assert dynamic_tags["dynamic-tag-1"]["created_at"] is not None + assert dynamic_tags["dynamic-tag-1"]["updated_at"] is not None + + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_list_tags_no_dynamic_tags(): + """ + Test list_tags when there are no dynamic tags in the spend table. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + + stored_tag = Mock() + stored_tag.tag_name = "stored-tag" + stored_tag.description = "A stored tag" + stored_tag.models = [] + stored_tag.model_info = None + stored_tag.spend = 0.0 + stored_tag.budget_id = None + stored_tag.created_at = datetime(2025, 1, 1) + stored_tag.updated_at = datetime(2025, 1, 1) + stored_tag.created_by = "user-123" + stored_tag.litellm_budget_table = None + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[stored_tag]) + + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[]) + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get("/tag/list", headers=headers) + + assert response.status_code == 200 + result = response.json() + assert len(result) == 1 + assert result[0]["name"] == "stored-tag" + + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_get_deployments_by_model_id(): """