mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(spend_logs): store litellm_call_id and match it in request_id lookups
Success spend rows are keyed by the upstream provider response id, so the x-litellm-call-id response header value never found them. Add a nullable indexed litellm_call_id column to LiteLLM_SpendLogs, populate it at write time, and widen every request_id lookup surface (/spend/logs, /spend/logs/ui, request details, ownership check) to match either id.
This commit is contained in:
mateo-berri
parent
d1320404fe
commit
f1dea17be1
9 changed files with 165 additions and 22 deletions
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "litellm_call_id" TEXT;
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
|
||||
|
|
@ -656,12 +656,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -3644,6 +3644,7 @@ class SpendLogsPayload(TypedDict):
|
|||
session_id: str | None
|
||||
request_duration_ms: int | None
|
||||
status: Literal["success", "failure"]
|
||||
litellm_call_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class SpanAttributes(str, enum.Enum):
|
||||
|
|
|
|||
|
|
@ -656,12 +656,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -229,10 +229,24 @@ async def _find_spend_logs(
|
|||
return rows
|
||||
|
||||
|
||||
class _RequestIdEquals(TypedDict):
|
||||
request_id: ReadOnly[str]
|
||||
|
||||
|
||||
class _LitellmCallIdEquals(TypedDict):
|
||||
litellm_call_id: ReadOnly[str]
|
||||
|
||||
|
||||
def _request_id_or_call_id_clause(request_id: str) -> tuple[_RequestIdEquals, _LitellmCallIdEquals]:
|
||||
request_id_clause: Final[_RequestIdEquals] = {"request_id": request_id}
|
||||
call_id_clause: Final[_LitellmCallIdEquals] = {"litellm_call_id": request_id}
|
||||
return (request_id_clause, call_id_clause)
|
||||
|
||||
|
||||
async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None:
|
||||
"""Read the single spend log row identified by ``request_id``."""
|
||||
return await _spend_logs_table(prisma_client).find_unique(
|
||||
where={"request_id": request_id},
|
||||
"""Read the single spend log row identified by ``request_id`` or ``litellm_call_id``."""
|
||||
return await _spend_logs_table(prisma_client).find_first(
|
||||
where={"OR": _request_id_or_call_id_clause(request_id)},
|
||||
include=None,
|
||||
)
|
||||
|
||||
|
|
@ -2543,7 +2557,6 @@ async def ui_view_spend_logs(
|
|||
("team_id", "team_id"),
|
||||
('"user"', "user"),
|
||||
("api_key", "api_key"),
|
||||
("request_id", "request_id"),
|
||||
("model", "model"),
|
||||
("model_id", "model_id"),
|
||||
("model_group", "model_group"),
|
||||
|
|
@ -2555,6 +2568,12 @@ async def ui_view_spend_logs(
|
|||
sql_params.append(val)
|
||||
p += 1
|
||||
|
||||
request_id_filter: Final = where_conditions.get("request_id")
|
||||
if isinstance(request_id_filter, str):
|
||||
sql_conditions.append(f"(request_id = ${p} OR litellm_call_id = ${p})")
|
||||
sql_params.append(request_id_filter)
|
||||
p += 1
|
||||
|
||||
# Multi-team OR filter: (user = $X OR team_id = ANY($Y))
|
||||
if permitted_team_ids:
|
||||
or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))'
|
||||
|
|
@ -2662,6 +2681,7 @@ async def ui_view_spend_logs(
|
|||
cache_hit, cache_key, request_tags, team_id,
|
||||
organization_id, end_user, requester_ip_address,
|
||||
session_id, status, mcp_namespaced_tool_name, agent_id,
|
||||
litellm_call_id,
|
||||
COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE {" AND ".join(sql_conditions)}
|
||||
|
|
@ -2735,7 +2755,7 @@ def _hydrate_spend_log_metadata(rows: Sequence[Mapping[str, object]]) -> None:
|
|||
|
||||
|
||||
def _cold_storage_object_key_from_metadata(
|
||||
metadata: str | dict | None,
|
||||
metadata: str | Mapping[str, object] | None,
|
||||
) -> str | None:
|
||||
if isinstance(metadata, str):
|
||||
try:
|
||||
|
|
@ -2870,7 +2890,7 @@ async def ui_view_request_response_for_request_id(
|
|||
sql_query: Final = """
|
||||
SELECT messages, response, proxy_server_request, metadata
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = $1
|
||||
WHERE request_id = $1 OR litellm_call_id = $1
|
||||
LIMIT 1
|
||||
"""
|
||||
db_result: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none(
|
||||
|
|
@ -2989,7 +3009,7 @@ async def view_spend_logs(
|
|||
start_date_iso: Final = start_date_obj.isoformat()
|
||||
end_date_iso: Final = end_date_obj.isoformat()
|
||||
|
||||
filter_query: Final = {
|
||||
filter_query: Final[dict[str, object]] = {
|
||||
"startTime": {
|
||||
"gte": start_date_iso, # Greater than or equal to Start Date
|
||||
"lte": end_date_iso, # Less than or equal to End Date
|
||||
|
|
@ -3002,7 +3022,7 @@ async def view_spend_logs(
|
|||
else:
|
||||
filter_query["api_key"] = api_key
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
filter_query["request_id"] = request_id
|
||||
filter_query["OR"] = _request_id_or_call_id_clause(request_id)
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
filter_query["user"] = user_id
|
||||
|
||||
|
|
@ -3073,7 +3093,7 @@ async def view_spend_logs(
|
|||
return response
|
||||
|
||||
else:
|
||||
scoped_filter: Final[dict[str, str]] = {}
|
||||
scoped_filter: Final[dict[str, object]] = {}
|
||||
if api_key is not None and isinstance(api_key, str):
|
||||
if api_key.startswith("sk-"):
|
||||
hashed_token = prisma_client.hash_token(token=api_key)
|
||||
|
|
@ -3081,7 +3101,7 @@ async def view_spend_logs(
|
|||
hashed_token = api_key
|
||||
scoped_filter["api_key"] = hashed_token
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
scoped_filter["request_id"] = request_id
|
||||
scoped_filter["OR"] = _request_id_or_call_id_clause(request_id)
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
scoped_filter["user"] = user_id
|
||||
|
||||
|
|
|
|||
|
|
@ -565,6 +565,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
status=_get_status_for_spend_log(
|
||||
metadata=metadata,
|
||||
),
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -656,12 +656,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -98,7 +98,10 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond)
|
||||
status = re.fullmatch(r"status = \$(\d+)", cond)
|
||||
api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond)
|
||||
if gte:
|
||||
req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond)
|
||||
if req_or_call:
|
||||
where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1]
|
||||
elif gte:
|
||||
date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1])
|
||||
elif lte:
|
||||
date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1])
|
||||
|
|
@ -410,7 +413,7 @@ async def test_assert_user_can_view_request_id_rejects_both_users_none():
|
|||
team_id = None
|
||||
|
||||
class MockSpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
async def find_first(self, where=None, include=None):
|
||||
return MockRow()
|
||||
|
||||
class MockDB:
|
||||
|
|
@ -453,6 +456,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat
|
|||
|
||||
ignored_keys = [
|
||||
"request_id",
|
||||
"litellm_call_id",
|
||||
"metadata.litellm_call_id",
|
||||
"session_id",
|
||||
"startTime",
|
||||
|
|
@ -2162,7 +2166,10 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window(
|
|||
def filter_fn(where):
|
||||
captured["where"] = where
|
||||
rows = _filter_logs_by_date_range(mock_spend_logs, where)
|
||||
if where.get("request_id"):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))]
|
||||
elif where.get("request_id"):
|
||||
rows = [r for r in rows if r["request_id"] == where["request_id"]]
|
||||
return rows
|
||||
|
||||
|
|
@ -2192,9 +2199,82 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window(
|
|||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert data["data"][0]["request_id"] == "req-old"
|
||||
# Query dropped the time window and scoped solely by the primary key.
|
||||
# Query dropped the time window and scoped solely by the id lookup.
|
||||
assert "startTime" not in captured["where"]
|
||||
assert captured["where"]["request_id"] == "req-old"
|
||||
assert captured["where"]["request_id_or_call_id"] == "req-old"
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_request_id_lookup_matches_litellm_call_id(
|
||||
client, monkeypatch
|
||||
):
|
||||
"""
|
||||
LIT-6302: success rows are keyed by the upstream provider response id, so a
|
||||
lookup with the x-litellm-call-id response header value found nothing. The id
|
||||
lookup now matches request_id OR litellm_call_id, resolving the header value.
|
||||
"""
|
||||
today = datetime.datetime.now(timezone.utc)
|
||||
mock_spend_logs = [
|
||||
{
|
||||
"id": "log_provider_keyed",
|
||||
"request_id": "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm",
|
||||
"litellm_call_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": "team1",
|
||||
"spend": 0.05,
|
||||
"startTime": today.isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log_other",
|
||||
"request_id": "chatcmpl-other",
|
||||
"litellm_call_id": "11111111-2222-3333-4444-555555555555",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": "team1",
|
||||
"spend": 0.01,
|
||||
"startTime": today.isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
def filter_fn(where):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
return [
|
||||
r
|
||||
for r in mock_spend_logs
|
||||
if rid_either in (r["request_id"], r.get("litellm_call_id"))
|
||||
]
|
||||
if where.get("request_id"):
|
||||
return [
|
||||
r for r in mock_spend_logs if r["request_id"] == where["request_id"]
|
||||
]
|
||||
return list(mock_spend_logs)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn),
|
||||
)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert (
|
||||
data["data"][0]["request_id"] == "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm"
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -2254,7 +2334,7 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc
|
|||
team_id = None
|
||||
|
||||
class _SpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
async def find_first(self, where=None, include=None):
|
||||
return _ForeignRow()
|
||||
|
||||
class _DB:
|
||||
|
|
@ -2307,7 +2387,10 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
|||
def filter_fn(where):
|
||||
captured["where"] = where
|
||||
rows = _filter_logs_by_date_range(mock_spend_logs, where)
|
||||
if where.get("request_id"):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))]
|
||||
elif where.get("request_id"):
|
||||
rows = [r for r in rows if r["request_id"] == where["request_id"]]
|
||||
return rows
|
||||
|
||||
|
|
@ -2317,10 +2400,10 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
|||
user = "user_1"
|
||||
team_id = "team1"
|
||||
|
||||
async def _find_unique(where, include=None):
|
||||
async def _find_first(where=None, include=None):
|
||||
return _OwnedRow()
|
||||
|
||||
mock_prisma.db.find_unique = _find_unique
|
||||
mock_prisma.db.find_first = _find_first
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
# A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends.
|
||||
|
|
@ -2345,7 +2428,7 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
|||
assert data["total"] == 1
|
||||
assert data["data"][0]["request_id"] == "req-old"
|
||||
assert "startTime" not in captured["where"]
|
||||
assert captured["where"]["request_id"] == "req-old"
|
||||
assert captured["where"]["request_id_or_call_id"] == "req-old"
|
||||
assert "user" not in captured["where"]
|
||||
assert "OR" not in captured["where"]
|
||||
finally:
|
||||
|
|
@ -4435,7 +4518,7 @@ async def test_view_spend_logs_internal_user_combines_user_with_request_id(
|
|||
where = mock_client.db.captured_where
|
||||
assert where is not None
|
||||
assert where["user"] == "internal-user-2"
|
||||
assert where["request_id"] == "req-abc"
|
||||
assert where["OR"] == ({"request_id": "req-abc"}, {"litellm_call_id": "req-abc"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -4462,7 +4545,7 @@ async def test_view_spend_logs_non_date_range_combines_user_with_request_id(
|
|||
where = mock_client.db.captured_where
|
||||
assert where is not None
|
||||
assert where["user"] == "internal-user-3"
|
||||
assert where["request_id"] == "req-xyz"
|
||||
assert where["OR"] == ({"request_id": "req-xyz"}, {"litellm_call_id": "req-xyz"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
|
|||
|
|
@ -1071,6 +1071,33 @@ def test_get_logging_payload_includes_agent_id_from_kwargs():
|
|||
), f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'"
|
||||
|
||||
|
||||
def test_get_logging_payload_populates_litellm_call_id_alongside_provider_request_id():
|
||||
"""
|
||||
LIT-6302: request_id stays the provider response id, so clients holding the
|
||||
x-litellm-call-id header value could never find their row. The payload now
|
||||
also carries litellm_call_id as its own column for lookups by either id.
|
||||
"""
|
||||
call_id = "b980eea9-5cd9-4099-93cd-8291e46c76fd"
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_call_id": call_id,
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
|
||||
},
|
||||
response_obj=litellm.ModelResponse(
|
||||
id="chatcmpl-provider-id",
|
||||
choices=[],
|
||||
usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
),
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["request_id"] == "chatcmpl-provider-id"
|
||||
assert payload["litellm_call_id"] == call_id
|
||||
|
||||
|
||||
@patch("litellm.proxy.proxy_server.master_key", None)
|
||||
@patch("litellm.proxy.proxy_server.general_settings", {})
|
||||
def test_get_logging_payload_includes_overhead_in_spend_logs_metadata():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue