From 7be7cc993be900a291375f860162b08ad53f40e9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 21 Oct 2026 06:12:29 +0000 Subject: [PATCH] refactor(spend-logs): drop unused key alias ORM lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_management_endpoints.py | 21 ------ .../test_spend_management_endpoints.py | 74 +++---------------- 2 files changed, 11 insertions(+), 84 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 2e839579c8e..2d00be6c2c3 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2775,27 +2775,6 @@ async def ui_view_spend_logs( where_conditions["metadata"] = metadata_filters[0] else: where_conditions["AND"] = [{"metadata": filter_cond} for filter_cond in metadata_filters] - if key_alias is not None: - active_key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( - where={"key_alias": key_alias}, - ) - deleted_key_rows: Final = await prisma_client.db.litellm_deletedverificationtoken.find_many( - where={"key_alias": key_alias}, - ) - key_tokens: Final = list(dict.fromkeys(row.token for row in (*active_key_rows, *deleted_key_rows))) - key_alias_filter: Final = {"api_key": {"in": key_tokens}} - if where_conditions.get("api_key") is None: - where_conditions["api_key"] = key_alias_filter["api_key"] - else: - existing_and_conditions: Final = where_conditions.get("AND", []) - where_conditions["AND"] = [ - *( - existing_and_conditions - if isinstance(existing_and_conditions, list) - else [existing_and_conditions] - ), - key_alias_filter, - ] if end_user is not None: where_conditions["end_user"] = end_user diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index b8e8e408f99..50aaa927bd5 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4757,7 +4757,7 @@ async def test_ui_view_spend_logs_with_error_message(client): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): - """Test exact alias resolution while preserving API key and metadata filters""" + """Test error-code filtering when a key-alias filter is supplied""" mock_spend_logs = [ { "id": "log1", @@ -4773,13 +4773,13 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): { "id": "log2", "request_id": "req2", - "api_key": "sk-other-key", + "api_key": "sk-test-key", "user": "test_user_2", "team_id": "team1", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4", - "metadata": '{"user_api_key_alias": "test-key-2", "error_information": {"error_code": "500"}}', + "metadata": '{"user_api_key_alias": "test-key-2", "error_information": {"error_code": "404"}}', }, { "id": "log3", @@ -4796,49 +4796,27 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): def filter_by_error_code_and_key_alias(where): metadata_filter = where.get("metadata", {}) - has_alias_substring_filter = any( - condition.get("metadata", {}).get("path") == ["user_api_key_alias"] - for condition in where.get("AND", []) - ) - if has_alias_substring_filter: - return mock_spend_logs - if ( - metadata_filter.get("path") != ["error_information", "error_code"] - or str(metadata_filter.get("equals", "")).strip('"') != "500" - ): - return mock_spend_logs - api_key_filter = where.get("api_key") - if api_key_filter == "sk-other-key": - return [] - if api_key_filter in ("sk-test-key", None) or api_key_filter == { - "in": ["sk-test-key", "sk-deleted-key"] - }: + if metadata_filter.get("path") == ["error_information", "error_code"] and str( + metadata_filter.get("equals", "") + ).strip('"') == "500": return [mock_spend_logs[2]] return mock_spend_logs - query_observer = MagicMock() - prisma_mock = make_ui_spend_logs_mock_prisma( - mock_spend_logs, - filter_by_error_code_and_key_alias, - query_observer=query_observer, - ) - active_key_lookup = AsyncMock(return_value=[MagicMock(token="sk-test-key")]) - deleted_key_lookup = AsyncMock(return_value=[MagicMock(token="sk-deleted-key")]) - prisma_mock.db.litellm_verificationtoken = MagicMock(find_many=active_key_lookup) - prisma_mock.db.litellm_deletedverificationtoken = MagicMock(find_many=deleted_key_lookup) - app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" ) try: - with patch.object(ps, "prisma_client", prisma_mock): + with patch.object( + ps, + "prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code_and_key_alias), + ): start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", params={ - "api_key": "sk-test-key", "error_code": "500", "key_alias": "test-key-1", "start_date": start_date, @@ -4858,36 +4836,6 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): assert metadata["user_api_key_alias"] == "test-key-1" assert "error_information" in metadata assert metadata["error_information"]["error_code"] == "500" - - mismatch_response = client.get( - "/spend/logs/ui", - params={ - "api_key": "sk-other-key", - "error_code": "500", - "key_alias": "test-key-1", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, - ) - assert mismatch_response.status_code == 200 - assert mismatch_response.json()["total"] == 0 - assert mismatch_response.json()["data"] == [] - page_query_call = next( - call - for call in query_observer.call_args_list - if 'FROM "LiteLLM_SpendLogs"' in call.args[0] and "ORDER BY" in call.args[0] - ) - page_query: Final = page_query_call.args[0] - page_params: Final = page_query_call.args[1] - alias_subquery: Final = re.search( - r'api_key IN \(SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = \$(\d+) ' - r'UNION SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE key_alias = \$\1\)', - page_query, - ) - assert alias_subquery is not None - assert page_params[int(alias_subquery.group(1)) - 1] == "test-key-1" - assert page_params.count("test-key-1") == 1 finally: app.dependency_overrides.pop(ps.user_api_key_auth, None)