Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 17 additions & 13 deletions litellm/proxy/management_endpoints/tag_management_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
19 changes: 0 additions & 19 deletions litellm/types/tag_management.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
from datetime import datetime
from typing import Dict, List, Optional

from pydantic import BaseModel
Expand Down Expand Up @@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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():
"""
Expand Down
Loading