mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
refactor(spend-logs): drop unused key alias ORM lookup
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
2e7e08ac33
commit
7be7cc993b
2 changed files with 11 additions and 84 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue