From beb51b743cffb99b3cdc59380efd153799d819ac Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 23 Jun 2026 15:29:51 -0700 Subject: [PATCH] Revert "fix(spend): surface request tags stored as JSON scalar strings in /spend/tags" This reverts commit 5e0c249210274ad65d48c6203ec646315ed35546. --- .../spend_management_endpoints.py | 30 ++--- .../test_spend_management_endpoints.py | 115 ------------------ 2 files changed, 13 insertions(+), 132 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 2ad86e5202e..0ba77dcd2f0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -218,6 +218,15 @@ async def view_spend_tags( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + # run the following SQL query on prisma + """ + SELECT + jsonb_array_elements_text(request_tags) AS individual_request_tag, + COUNT(*) AS log_count, + SUM(spend) AS total_spend + FROM "LiteLLM_SpendLogs" + GROUP BY individual_request_tag; + """ response = await get_spend_by_tags( start_date=start_date, end_date=end_date, prisma_client=prisma_client ) @@ -3239,24 +3248,11 @@ async def get_spend_by_tags( prisma_client: PrismaClient, start_date=None, end_date=None ): response = await prisma_client.db.query_raw(""" - WITH normalized_tags AS ( - SELECT - spend, - CASE - WHEN jsonb_typeof(request_tags) = 'array' THEN request_tags - WHEN jsonb_typeof(request_tags) = 'string' - AND (request_tags #>> '{}') ~ '^\\s*\\[' - THEN (request_tags #>> '{}')::jsonb - ELSE NULL - END AS tags - FROM "LiteLLM_SpendLogs" - ) SELECT - jsonb_array_elements_text(tags) AS individual_request_tag, - COUNT(*) AS log_count, - SUM(spend) AS total_spend - FROM normalized_tags - WHERE jsonb_typeof(tags) = 'array' + jsonb_array_elements_text(request_tags) AS individual_request_tag, + COUNT(*) AS log_count, + SUM(spend) AS total_spend + FROM "LiteLLM_SpendLogs" GROUP BY individual_request_tag; """) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index b708827d4b8..0b583129591 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -97,10 +97,6 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return MockPrismaClient() -from collections import Counter -from typing import Dict, List, Optional, Tuple - -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( LitellmUserRoles, Member, @@ -3790,114 +3786,3 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict assert body["data"][0]["metadata"] == {} finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) - - -class _ScalarExtractionError(Exception): - """Mirrors PostgreSQL's "cannot extract elements from a scalar" error that - ``jsonb_array_elements_text`` raises when a row's request_tags is not a JSON - array. A single such row aborts the whole GROUP BY, which is the LIT-3906 - bug.""" - - -def _jsonb_typeof(literal: str) -> str: - value = json.loads(literal) - if isinstance(value, list): - return "array" - if isinstance(value, str): - return "string" - if isinstance(value, dict): - return "object" - if value is None: - return "null" - return "number" - - -def _normalized_tags(literal: str, query_guards_strings: bool) -> Optional[List[str]]: - typeof = _jsonb_typeof(literal) - if typeof == "array": - return json.loads(literal) - if not query_guards_strings: - raise _ScalarExtractionError("cannot extract elements from a scalar") - if typeof != "string": - return None - inner = json.loads(literal) - if not inner.lstrip().startswith("["): - return None - parsed = json.loads(inner) - return parsed if isinstance(parsed, list) else None - - -class _FakeSpendLogsJsonbTable: - """Faithful emulator of the subset of PostgreSQL JSONB behaviour that - ``get_spend_by_tags`` relies on. - - Each seeded row stores ``request_tags`` exactly as the spend-log write path - produces it: the tag list is ``safe_dumps``'d to a JSON string and Prisma's - ``create_many`` JSON-encodes that string again into the ``Json`` column, so - the column literal is ``json.dumps(json.dumps(tags))`` - a JSON scalar - string, not a JSON array. ``jsonb_array_elements_text`` raises on a scalar, - which is the LIT-3906 bug. - """ - - def __init__(self, rows: Tuple[Tuple[str, float], ...]) -> None: - self._rows = rows - - async def query_raw(self, query: str, *args: object) -> List[Dict[str, object]]: - guards_strings = "jsonb_typeof(request_tags) = 'string'" in query - expanded = tuple( - (tag, spend) - for literal, spend in self._rows - for tag in (_normalized_tags(literal, guards_strings) or ()) - ) - counts = Counter(tag for tag, _ in expanded) - spends: Dict[str, float] = { - tag: sum(spend for t, spend in expanded if t == tag) for tag in counts - } - return [ - { - "individual_request_tag": tag, - "log_count": counts[tag], - "total_spend": spends[tag], - } - for tag in counts - ] - - -def _double_encoded_request_tags(tags: List[str]) -> str: - return json.dumps(safe_dumps(tags)) - - -@pytest.mark.asyncio -async def test_get_spend_by_tags_surfaces_scalar_string_tags() -> None: - """Regression for LIT-3906: tags stored as a JSON scalar string (the shape - the write path produces) must still surface from /spend/tags with the - summed spend and correct log count.""" - rows = ( - (_double_encoded_request_tags(["prod-tag", "shared"]), 0.10), - (_double_encoded_request_tags(["prod-tag"]), 0.05), - (_double_encoded_request_tags(["shared"]), 0.02), - (json.dumps(["already-array"]), 0.03), - (json.dumps(None), 0.99), - (json.dumps("not-json-at-all"), 0.99), - ) - - prisma_client = MagicMock() - prisma_client.db = _FakeSpendLogsJsonbTable(rows) - - response = await spend_management_endpoints.get_spend_by_tags( - prisma_client=prisma_client - ) - - by_tag = {r["individual_request_tag"]: r for r in response} - - assert "prod-tag" in by_tag, "double-encoded tag never surfaced from /spend/tags" - assert by_tag["prod-tag"]["log_count"] == 2 - assert by_tag["prod-tag"]["total_spend"] == pytest.approx(0.15) - - assert by_tag["shared"]["log_count"] == 2 - assert by_tag["shared"]["total_spend"] == pytest.approx(0.12) - - assert by_tag["already-array"]["log_count"] == 1 - assert by_tag["already-array"]["total_spend"] == pytest.approx(0.03) - - assert "not-json-at-all" not in by_tag