scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

feat(focus): include per-user request tags in Vantage export

Browse source 
Source request tags from LiteLLM_SpendLogs (which carries both user_id and
request_tags per request) so exported FOCUS Tags include request-level tags,
attributed to the exact user even when an api_key is shared. Scope the
SpendLogs aggregation to the export time window via the startTime index.
This commit is contained in:
yucheng-berriai 2026-06-22 23:30:14 -07:00
parent c546b58c09
commit 18eac37a4a
4 changed files with 272 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/integrations/focus/database.py
View file

@ -32,14 +32,21 @@ class FocusLiteLLMDatabase:
client = self._ensure_prisma_client()
where_clauses: list[str] = []
tag_time_clauses: list[str] = []
query_params: list[Any] = []
placeholder_index = 1
if start_time_utc:
where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz")
tag_time_clauses.append(
f"sl.\"startTime\" >= (${placeholder_index}::timestamptz AT TIME ZONE 'UTC')"
)
query_params.append(start_time_utc)
placeholder_index += 1
if end_time_utc:
where_clauses.append(f"dus.updated_at <= ${placeholder_index}::timestamptz")
tag_time_clauses.append(
f"sl.\"startTime\" <= (${placeholder_index}::timestamptz AT TIME ZONE 'UTC')"
)
query_params.append(end_time_utc)
placeholder_index += 1
@ -47,6 +54,10 @@ class FocusLiteLLMDatabase:
if where_clauses:
where_clause = "WHERE " + " AND ".join(where_clauses)
tag_time_clause = ""
if tag_time_clauses:
tag_time_clause = "WHERE " + " AND ".join(tag_time_clauses)
limit_clause = ""
if limit is not None:
try:
@ -82,13 +93,42 @@ class FocusLiteLLMDatabase:
tt.team_alias,
ut.user_email as user_email,
COALESCE(vt.organization_id, tt.organization_id) as organization_id,
ot.organization_alias as organization_alias
ot.organization_alias as organization_alias,
tag_rollup.request_tags as request_tags
FROM "LiteLLM_DailyUserSpend" dus
LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
LEFT JOIN "LiteLLM_OrganizationTable" ot
ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id)
LEFT JOIN (
SELECT
sl."user" AS user_id,
to_char(sl."startTime", 'YYYY-MM-DD') AS date,
sl.api_key,
sl.model,
sl.custom_llm_provider,
sl.mcp_namespaced_tool_name,
ARRAY_AGG(DISTINCT tag_value ORDER BY tag_value) FILTER (WHERE tag_value IS NOT NULL) AS request_tags
FROM "LiteLLM_SpendLogs" sl
LEFT JOIN LATERAL jsonb_array_elements_text(
CASE
WHEN jsonb_typeof(sl.request_tags::jsonb) = 'array'
THEN sl.request_tags::jsonb
ELSE '[]'::jsonb
END
) AS elem(tag_value) ON TRUE
{tag_time_clause}
GROUP BY sl."user",
to_char(sl."startTime", 'YYYY-MM-DD'),
sl.api_key, sl.model, sl.custom_llm_provider, sl.mcp_namespaced_tool_name
) tag_rollup
ON COALESCE(tag_rollup.user_id, '') = COALESCE(dus.user_id, '')
AND tag_rollup.date = dus.date
AND tag_rollup.api_key = dus.api_key
AND COALESCE(tag_rollup.model, '') = COALESCE(dus.model, '')
AND COALESCE(tag_rollup.custom_llm_provider, '') = COALESCE(dus.custom_llm_provider, '')
AND COALESCE(tag_rollup.mcp_namespaced_tool_name, '') = COALESCE(dus.mcp_namespaced_tool_name, '')
{where_clause}
ORDER BY dus.date DESC, dus.created_at DESC
{limit_clause}

53
litellm/integrations/focus/transformer.py
View file

@ -22,22 +22,59 @@ _TAG_KEYS = (
"custom_llm_provider",
)
_REQUEST_TAGS_KEY = "request_tags"
_REQUEST_TAGS_TRUNCATED_KEY = "request_tags_truncated"
def _build_tags_expr(available_keys: list[str]) -> pl.Expr:
# Request tags are caller-controlled and unbounded. Cap how many (and how long)
# land in a single row's Tags blob so it cannot exceed a destination's per-row
# size limit (Vantage drops any row over 2 MB). The cap is on characters; even
# the worst case (64 tags x 128 multibyte chars, double-JSON-encoded then CSV-
# escaped) is on the order of ~115 KB, well under 2 MB.
_MAX_REQUEST_TAGS = 64
_MAX_TAG_LENGTH = 128
def _build_tags_expr(tag_columns: list[str]) -> pl.Expr:
"""Build a Polars expression that produces a JSON Tags string per row.
Uses ``pl.struct`` + ``map_elements`` to avoid materialising the entire
DataFrame to a list of Python dicts. The JSON serialisation callback
still runs in Python (GIL-bound), but struct-packing and loop dispatch
are handled by Polars' Rust engine.
Request-level tags arrive as a list column and are encoded as a JSON
array string so the Tags map stays a flat string-to-string object. The
list is capped to a bounded count/length so a flood of caller-supplied
tags cannot push the row past a destination's per-row size limit; when
that happens a ``request_tags_truncated`` marker carries the true count.
"""
def _struct_to_json(row: dict) -> str:
tags = {k: str(v) for k, v in row.items() if v is not None}
return json.dumps(tags) if tags else "{}"
def _struct_to_json(row: dict[str, object]) -> str:
metadata = {
k: str(v)
for k, v in row.items()
if k != _REQUEST_TAGS_KEY and v is not None
}
raw_tags = row.get(_REQUEST_TAGS_KEY)
if not raw_tags or not isinstance(raw_tags, (list, tuple)):
return json.dumps(metadata) if metadata else "{}"
total = len(raw_tags)
capped: list[str] = [
str(t)[:_MAX_TAG_LENGTH] for t in raw_tags[:_MAX_REQUEST_TAGS]
]
tags = {
**metadata,
_REQUEST_TAGS_KEY: json.dumps(capped),
**(
{_REQUEST_TAGS_TRUNCATED_KEY: str(total)}
if total > _MAX_REQUEST_TAGS
else {}
),
}
return json.dumps(tags)
return (
pl.struct(available_keys)
pl.struct(tag_columns)
.map_elements(_struct_to_json, return_dtype=pl.String)
.alias("Tags")
)
@ -54,9 +91,9 @@ class FocusTransformer:
return pl.DataFrame(schema=self.schema)
# Build Tags JSON from metadata columns using vectorized Polars expression
available_keys = [k for k in _TAG_KEYS if k in frame.columns]
if available_keys:
frame = frame.with_columns(_build_tags_expr(available_keys))
tag_columns = [k for k in (*_TAG_KEYS, _REQUEST_TAGS_KEY) if k in frame.columns]
if tag_columns:
frame = frame.with_columns(_build_tags_expr(tag_columns))
else:
frame = frame.with_columns(pl.lit("{}").alias("Tags"))

79
tests/test_litellm/integrations/focus/test_focus_database.py
View file

@ -1,5 +1,6 @@
"""Tests for FocusLiteLLMDatabase query construction."""
import json
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
@ -7,6 +8,7 @@ from unittest.mock import AsyncMock
import pytest
from litellm.integrations.focus.database import FocusLiteLLMDatabase
from litellm.integrations.focus.transformer import FocusTransformer
def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return):
@ -45,8 +47,12 @@ async def test_should_execute_without_filters(monkeypatch: pytest.MonkeyPatch):
result = await db.get_usage_data()
query_text, *params = query_mock.await_args.args
assert "WHERE" not in query_text
assert "dus.updated_at >=" not in query_text
assert "dus.updated_at <=" not in query_text
assert "LIMIT $" not in query_text
# no stray WHERE clause anywhere when no window is given; the only legitimate
# WHERE is the ARRAY_AGG FILTER, so strip that before asserting
assert "WHERE" not in query_text.replace("FILTER (WHERE", "")
assert params == []
assert result.height == 1
assert result["id"][0] == 1
@ -87,3 +93,74 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
)
assert "ot.organization_alias as organization_alias" in query_text
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
@pytest.mark.asyncio
async def test_should_join_spend_logs_for_per_user_request_tags(
monkeypatch: pytest.MonkeyPatch,
):
db, query_mock = _setup_db(monkeypatch, [])
await db.get_usage_data()
query_text, *_ = query_mock.await_args.args
assert 'FROM "LiteLLM_SpendLogs" sl' in query_text
assert "jsonb_array_elements_text" in query_text
assert "tag_rollup.request_tags as request_tags" in query_text
# exact per-user attribution: tags join on user_id so a shared api_key
# never leaks one user's tags onto another user's export row
assert "COALESCE(tag_rollup.user_id, '') = COALESCE(dus.user_id, '')" in query_text
@pytest.mark.asyncio
async def test_should_scope_tag_subquery_to_time_window(
monkeypatch: pytest.MonkeyPatch,
):
"""The SpendLogs tag aggregation must be bounded to the requested time
window (using the startTime index) instead of scanning all history."""
start = datetime(2024, 1, 1, tzinfo=timezone.utc)
end = datetime(2024, 1, 2, tzinfo=timezone.utc)
db, query_mock = _setup_db(monkeypatch, [])
await db.get_usage_data(start_time_utc=start, end_time_utc=end)
query_text, *params = query_mock.await_args.args
# window filter converts the bound param to the column's naive-UTC frame so
# it is independent of the Postgres session TimeZone
assert "sl.\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')" in query_text
assert "sl.\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')" in query_text
# bucketing must NOT apply AT TIME ZONE to the column (startTime is already
# naive-UTC; rendering it in the session TZ would shift the date and break
# the join with dus.date on non-UTC sessions)
assert "to_char(sl.\"startTime\", 'YYYY-MM-DD')" in query_text
assert "AT TIME ZONE 'UTC', 'YYYY-MM-DD'" not in query_text
assert params == [start, end]
@pytest.mark.asyncio
async def test_request_tags_flow_from_db_array_into_focus_tags(
monkeypatch: pytest.MonkeyPatch,
):
"""Postgres returns the aggregated tags as an array; it must survive the
DataFrame round-trip and land inside the transformed FOCUS Tags column."""
row = {
"date": "2024-01-01",
"user_id": "user",
"api_key": "sk-test",
"api_key_alias": "prod-key",
"model": "gpt-4o",
"model_group": "gpt-4o",
"custom_llm_provider": "openai",
"spend": 0.05,
"api_requests": 1,
"team_id": "team-1",
"team_alias": "Platform",
"request_tags": ["prod", "checkout"],
}
db, _ = _setup_db(monkeypatch, [row])
frame = await db.get_usage_data()
normalized = FocusTransformer().transform(frame)
tags = json.loads(normalized["Tags"][0])
assert json.loads(tags["request_tags"]) == ["prod", "checkout"]

108
tests/test_litellm/integrations/focus/test_transformer.py
View file

@ -38,6 +38,114 @@ def test_should_include_organization_fields_in_tags():
assert tags["team_id"] == "team-1"
def test_should_include_request_tags_in_tags():
frame = pl.DataFrame(
{
"date": [date(2024, 1, 2)],
"spend": [1.25],
"api_requests": [1],
"api_key": ["hashed-key"],
"api_key_alias": ["prod-key"],
"model": ["gpt-4o"],
"model_group": ["gpt-4o"],
"custom_llm_provider": ["openai"],
"team_id": ["team-1"],
"team_alias": ["Platform"],
"request_tags": [["prod", "checkout"]],
}
)
normalized = FocusTransformer().transform(frame)
tags = json.loads(normalized["Tags"][0])
assert json.loads(tags["request_tags"]) == ["prod", "checkout"]
assert tags["team_id"] == "team-1"
assert "request_tags_truncated" not in tags
def test_should_cap_unbounded_request_tags():
"""Caller-supplied tags are unbounded; the export must cap how many and how
long they are so one row cannot exceed a destination's per-row size limit,
and must record the true count via request_tags_truncated."""
long_tag = "x" * 500
many_tags = [f"tag-{i}" for i in range(70)]
frame = pl.DataFrame(
{
"date": [date(2024, 1, 2)],
"spend": [1.0],
"api_requests": [1],
"api_key": ["hashed-key"],
"api_key_alias": ["prod-key"],
"model": ["gpt-4o"],
"model_group": ["gpt-4o"],
"custom_llm_provider": ["openai"],
"team_id": ["team-1"],
"team_alias": ["Platform"],
"request_tags": [[long_tag, *many_tags]],
}
)
normalized = FocusTransformer().transform(frame)
tags = json.loads(normalized["Tags"][0])
emitted = json.loads(tags["request_tags"])
assert len(emitted) == 64 # _MAX_REQUEST_TAGS
assert all(len(t) <= 128 for t in emitted) # _MAX_TAG_LENGTH
assert tags["request_tags_truncated"] == "71" # 1 long + 70
def test_should_not_mark_truncated_at_exactly_the_cap():
"""Exactly _MAX_REQUEST_TAGS tags is not truncation: all are kept and no
marker is emitted (guards the `>` vs `>=` boundary)."""
exactly_cap = [f"tag-{i}" for i in range(64)]
frame = pl.DataFrame(
{
"date": [date(2024, 1, 2)],
"spend": [1.0],
"api_requests": [1],
"api_key": ["hashed-key"],
"api_key_alias": ["prod-key"],
"model": ["gpt-4o"],
"model_group": ["gpt-4o"],
"custom_llm_provider": ["openai"],
"team_id": ["team-1"],
"team_alias": ["Platform"],
"request_tags": [exactly_cap],
}
)
normalized = FocusTransformer().transform(frame)
tags = json.loads(normalized["Tags"][0])
assert len(json.loads(tags["request_tags"])) == 64
assert "request_tags_truncated" not in tags
def test_should_omit_request_tags_when_empty():
frame = pl.DataFrame(
{
"date": [date(2024, 1, 2)],
"spend": [1.25],
"api_requests": [1],
"api_key": ["hashed-key"],
"api_key_alias": ["prod-key"],
"model": ["gpt-4o"],
"model_group": ["gpt-4o"],
"custom_llm_provider": ["openai"],
"team_id": ["team-1"],
"team_alias": ["Platform"],
}
).with_columns(
pl.Series("request_tags", [None], dtype=pl.List(pl.String)),
)
normalized = FocusTransformer().transform(frame)
tags = json.loads(normalized["Tags"][0])
assert "request_tags" not in tags
assert tags["team_id"] == "team-1"
def test_should_omit_missing_organization_fields_from_tags():
frame = pl.DataFrame(
{