mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(spend-logs): match key aliases exactly
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
7ae46f2768
commit
2e7e08ac33
3 changed files with 221 additions and 32 deletions
|
|
@ -2575,7 +2575,7 @@ async def ui_view_spend_logs(
|
|||
description="Filter logs by model ID (litellm model deployment id)",
|
||||
),
|
||||
model_group: str | None = fastapi.Query(default=None, description="Filter logs by model group"),
|
||||
key_alias: str | None = fastapi.Query(default=None, description="Filter logs by key alias"),
|
||||
key_alias: str | None = fastapi.Query(default=None, description="Filter logs by exact key alias match"),
|
||||
end_user: str | None = fastapi.Query(default=None, description="Filter logs by end user"),
|
||||
error_code: str | None = fastapi.Query(default=None, description="Filter logs by error code (e.g., '404', '500')"),
|
||||
error_message: str | None = fastapi.Query(
|
||||
|
|
@ -2754,14 +2754,6 @@ async def ui_view_spend_logs(
|
|||
|
||||
# Build metadata filters
|
||||
metadata_filters: Final = []
|
||||
if key_alias is not None:
|
||||
metadata_filters.append(
|
||||
{
|
||||
"path": ["user_api_key_alias"],
|
||||
"string_contains": key_alias,
|
||||
}
|
||||
)
|
||||
|
||||
if error_code is not None:
|
||||
metadata_filters.append(
|
||||
{
|
||||
|
|
@ -2783,6 +2775,27 @@ 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
|
||||
|
||||
|
|
@ -2929,8 +2942,11 @@ async def ui_view_spend_logs(
|
|||
|
||||
# Metadata JSON filters (PostgreSQL JSONB operators)
|
||||
if key_alias is not None:
|
||||
sql_conditions.append(f"metadata->>'user_api_key_alias' LIKE ${p}")
|
||||
sql_params.append(f"%{key_alias}%")
|
||||
sql_conditions.append(
|
||||
f'api_key IN (SELECT token FROM "LiteLLM_VerificationToken" WHERE key_alias = ${p} '
|
||||
f'UNION SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE key_alias = ${p})'
|
||||
)
|
||||
sql_params.append(key_alias)
|
||||
p += 1
|
||||
if error_code is not None:
|
||||
sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}")
|
||||
|
|
|
|||
|
|
@ -4,14 +4,17 @@ from collections.abc import AsyncIterator
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import httpx
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.client import Gateway, delete_key_if_present, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from prisma import Prisma
|
||||
from psycopg import sql
|
||||
from psycopg.types.json import Jsonb
|
||||
|
|
@ -77,6 +80,37 @@ def _seed_rows(
|
|||
)
|
||||
|
||||
|
||||
def _seed_alias_rows(
|
||||
connection: psycopg.Connection,
|
||||
rows: tuple[tuple[str, str, str | None, str], ...],
|
||||
session_id: str,
|
||||
started: datetime,
|
||||
) -> None:
|
||||
utc_timestamp: Final = started.astimezone(timezone.utc).replace(tzinfo=None)
|
||||
with connection.cursor() as cursor:
|
||||
cursor.executemany(
|
||||
sql.SQL(
|
||||
'INSERT INTO {} (request_id, api_key, "user", session_id, "startTime", "endTime", '
|
||||
"metadata, messages, response, call_type) "
|
||||
"VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, 'acompletion')"
|
||||
).format(sql.Identifier("public", "LiteLLM_SpendLogs")),
|
||||
tuple(
|
||||
(
|
||||
request_id,
|
||||
api_key,
|
||||
user_id,
|
||||
session_id,
|
||||
utc_timestamp,
|
||||
utc_timestamp,
|
||||
Jsonb({"user_api_key_alias": key_alias}),
|
||||
Jsonb([{"role": "user", "content": request_id + " payload"}]),
|
||||
Jsonb({"id": request_id}),
|
||||
)
|
||||
for request_id, api_key, user_id, key_alias in rows
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def spend_database() -> AsyncIterator[Prisma]:
|
||||
schema: Final = f"integration_spend_scope_{uuid.uuid4().hex}"
|
||||
|
|
@ -157,8 +191,8 @@ async def test_payload_sql_filters_foreign_collisions_and_prefers_exact_ids_for_
|
|||
assert tuple(row.request_id for row in REQUEST_IDS.validate_python(result)) == expected
|
||||
|
||||
|
||||
def _delete_session(session_id: str) -> None:
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
||||
def _delete_session(session_id: str, database_url: str | None = None) -> None:
|
||||
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:
|
||||
connection.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE session_id = %s', (session_id,))
|
||||
|
||||
|
||||
|
|
@ -240,6 +274,99 @@ def _chat_request_id(gateway: Gateway, model: str, key: str) -> str:
|
|||
return str(body["id"])
|
||||
|
||||
|
||||
def test_spend_log_key_alias_is_exact_includes_deleted_keys_and_preserves_scope(tmp_path: Path) -> None:
|
||||
with httpx.Client(base_url="http://127.0.0.1:1", timeout=15, trust_env=False) as source_client:
|
||||
source_gateway: Final = Gateway(source_client, "sk-integration-master", "http://127.0.0.1:8190")
|
||||
with owned_proxy_process(source_gateway, tmp_path, {}) as owned:
|
||||
gateway: Final = owned.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
assert gateway.database_url is not None
|
||||
database_url: Final = gateway.database_url
|
||||
owner: Final = scenario.user(user_role="internal_user")
|
||||
other_owner: Final = scenario.user(user_role="internal_user")
|
||||
key_alias: Final = f"prod-key-{uuid.uuid4().hex}"
|
||||
overlapping_key_alias: Final = f"{key_alias}-2"
|
||||
deleted_key_alias: Final = f"deleted-key-{uuid.uuid4().hex}"
|
||||
key: Final = scenario.key(user_id=owner, key_alias=key_alias)
|
||||
overlapping_key: Final = scenario.key(
|
||||
user_id=other_owner,
|
||||
key_alias=overlapping_key_alias,
|
||||
)
|
||||
deleted_key_response: Final = gateway.post(
|
||||
"/key/generate",
|
||||
{"key_alias": deleted_key_alias},
|
||||
)
|
||||
deleted_key: Final = string_value(deleted_key_response["key"])
|
||||
scenario.cleanups.callback(delete_key_if_present, gateway, deleted_key)
|
||||
request_ids: Final = tuple(f"alias-{uuid.uuid4().hex}" for _ in range(3))
|
||||
session_id: Final = f"key-alias-{uuid.uuid4().hex}"
|
||||
started: Final = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
scenario.cleanups.callback(_delete_session, session_id, database_url)
|
||||
with psycopg.connect(database_url) as connection:
|
||||
_seed_alias_rows(
|
||||
connection,
|
||||
(
|
||||
(request_ids[0], sha256(key.encode()).hexdigest(), owner, key_alias),
|
||||
(
|
||||
request_ids[1],
|
||||
sha256(overlapping_key.encode()).hexdigest(),
|
||||
other_owner,
|
||||
overlapping_key_alias,
|
||||
),
|
||||
(
|
||||
request_ids[2],
|
||||
sha256(deleted_key.encode()).hexdigest(),
|
||||
other_owner,
|
||||
deleted_key_alias,
|
||||
),
|
||||
),
|
||||
session_id,
|
||||
started,
|
||||
)
|
||||
gateway.post("/key/delete", {"keys": [deleted_key]})
|
||||
date_filters: Final = {
|
||||
"start_date": started.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"end_date": (started + timedelta(hours=2)).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
}
|
||||
|
||||
for path in ("/spend/logs/ui", "/spend/logs/v2"):
|
||||
exact: Final = gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={**date_filters, "key_alias": key_alias},
|
||||
)
|
||||
assert exact.status_code == 200, exact.text
|
||||
assert exact.json()["total"] == 1, exact.text
|
||||
assert [row["request_id"] for row in exact.json()["data"]] == [request_ids[0]], exact.text
|
||||
|
||||
deleted: Final = gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={**date_filters, "key_alias": deleted_key_alias},
|
||||
)
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert deleted.json()["total"] == 1, deleted.text
|
||||
assert [row["request_id"] for row in deleted.json()["data"]] == [request_ids[2]], deleted.text
|
||||
|
||||
absent: Final = gateway.request(
|
||||
"GET",
|
||||
path,
|
||||
params={**date_filters, "key_alias": f"absent-{uuid.uuid4().hex}"},
|
||||
)
|
||||
assert absent.status_code == 200, absent.text
|
||||
assert absent.json()["total"] == 0, absent.text
|
||||
assert absent.json()["data"] == [], absent.text
|
||||
|
||||
foreign_alias: Final = gateway.request(
|
||||
"GET",
|
||||
"/spend/logs/ui",
|
||||
params={**date_filters, "key_alias": overlapping_key_alias},
|
||||
key=key,
|
||||
)
|
||||
assert foreign_alias.status_code == 200, foreign_alias.text
|
||||
assert foreign_alias.json()["data"] == [], foreign_alias.text
|
||||
|
||||
|
||||
def _legacy_spend_rows(gateway: Gateway, key: str, params: dict[str, str]) -> list[dict[str, object]]:
|
||||
response: Final = gateway.request("GET", "/spend/logs", key=key, params=params)
|
||||
assert response.status_code == 200, response.text
|
||||
|
|
|
|||
|
|
@ -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 merging error_code and key_alias filters with AND logic"""
|
||||
"""Test exact alias resolution while preserving API key and metadata filters"""
|
||||
mock_spend_logs = [
|
||||
{
|
||||
"id": "log1",
|
||||
|
|
@ -4773,7 +4773,7 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client):
|
|||
{
|
||||
"id": "log2",
|
||||
"request_id": "req2",
|
||||
"api_key": "sk-test-key",
|
||||
"api_key": "sk-other-key",
|
||||
"user": "test_user_2",
|
||||
"team_id": "team1",
|
||||
"spend": 0.10,
|
||||
|
|
@ -4795,34 +4795,50 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client):
|
|||
]
|
||||
|
||||
def filter_by_error_code_and_key_alias(where):
|
||||
if "AND" in where:
|
||||
key_alias = error_code = None
|
||||
for cond in where["AND"]:
|
||||
if "metadata" in cond:
|
||||
mf = cond["metadata"]
|
||||
if mf.get("path") == ["user_api_key_alias"]:
|
||||
key_alias = mf.get("string_contains")
|
||||
elif mf.get("path") == ["error_information", "error_code"]:
|
||||
error_code = str(mf.get("equals", "")).strip('"')
|
||||
if key_alias == "test-key-1" and error_code == "500":
|
||||
return [mock_spend_logs[2]]
|
||||
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"]
|
||||
}:
|
||||
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",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code_and_key_alias),
|
||||
):
|
||||
with patch.object(ps, "prisma_client", prisma_mock):
|
||||
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,
|
||||
|
|
@ -4842,6 +4858,36 @@ 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