diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 3ae3f6b53ac..03042774029 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -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} diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index a17df29b912..c314085123f 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -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")) diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index d77af2dd170..4aead37ebdc 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -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"] diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/test_litellm/integrations/focus/test_transformer.py index 4461d19efde..5748905d8b6 100644 --- a/tests/test_litellm/integrations/focus/test_transformer.py +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -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( {