feat(vantage): include organization metadata in FOCUS Tags export (#28184)
* feat(vantage): include organization metadata in FOCUS Tags export Join LiteLLM_OrganizationTable when building Vantage/FOCUS export rows so organization_id and organization_alias appear in Tags for org-level filtering. Co-authored-by: Cursor <cursoragent@cursor.com> * test(focus): include api_requests in organization Tags tests FocusTransformer now requires api_requests after staging merge; add the column to test fixtures so integrations CI can run the Tags assertions. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
c24a3603d9
commit
411bd3da5b
@ -80,11 +80,15 @@ class FocusLiteLLMDatabase:
|
||||
vt.team_id,
|
||||
vt.key_alias as api_key_alias,
|
||||
tt.team_alias,
|
||||
ut.user_email as user_email
|
||||
ut.user_email as user_email,
|
||||
COALESCE(vt.organization_id, tt.organization_id) as organization_id,
|
||||
ot.organization_alias as organization_alias
|
||||
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)
|
||||
{where_clause}
|
||||
ORDER BY dus.date DESC, dus.created_at DESC
|
||||
{limit_clause}
|
||||
|
||||
@ -12,6 +12,8 @@ from .schema import FOCUS_NORMALIZED_SCHEMA
|
||||
_TAG_KEYS = (
|
||||
"team_id",
|
||||
"team_alias",
|
||||
"organization_id",
|
||||
"organization_alias",
|
||||
"user_id",
|
||||
"user_email",
|
||||
"api_key_alias",
|
||||
|
||||
@ -72,3 +72,18 @@ async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch):
|
||||
await db.get_usage_data(limit="invalid")
|
||||
|
||||
assert query_mock.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
|
||||
db, query_mock = _setup_db(monkeypatch, [])
|
||||
|
||||
await db.get_usage_data()
|
||||
|
||||
query_text, *_ = query_mock.await_args.args
|
||||
assert (
|
||||
"COALESCE(vt.organization_id, tt.organization_id) as organization_id"
|
||||
in query_text
|
||||
)
|
||||
assert "ot.organization_alias as organization_alias" in query_text
|
||||
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
|
||||
|
||||
62
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
62
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
@ -0,0 +1,62 @@
|
||||
"""Tests for FocusTransformer organization metadata in Tags."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import date
|
||||
|
||||
import polars as pl
|
||||
|
||||
from litellm.integrations.focus.transformer import FocusTransformer
|
||||
|
||||
|
||||
def test_should_include_organization_fields_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"],
|
||||
"organization_id": ["org-123"],
|
||||
"organization_alias": ["Acme Corp"],
|
||||
"user_id": ["user-1"],
|
||||
"user_email": ["user@example.com"],
|
||||
}
|
||||
)
|
||||
|
||||
normalized = FocusTransformer().transform(frame)
|
||||
|
||||
tags = json.loads(normalized["Tags"][0])
|
||||
assert tags["organization_id"] == "org-123"
|
||||
assert tags["organization_alias"] == "Acme Corp"
|
||||
assert tags["team_id"] == "team-1"
|
||||
|
||||
|
||||
def test_should_omit_missing_organization_fields_from_tags():
|
||||
frame = pl.DataFrame(
|
||||
{
|
||||
"date": [date(2024, 1, 2)],
|
||||
"spend": [0.5],
|
||||
"api_requests": [1],
|
||||
"api_key": ["hashed-key"],
|
||||
"api_key_alias": ["prod-key"],
|
||||
"model": ["gpt-4o-mini"],
|
||||
"model_group": ["gpt-4o-mini"],
|
||||
"custom_llm_provider": ["openai"],
|
||||
"team_id": ["team-1"],
|
||||
"team_alias": ["Platform"],
|
||||
}
|
||||
)
|
||||
|
||||
normalized = FocusTransformer().transform(frame)
|
||||
|
||||
tags = json.loads(normalized["Tags"][0])
|
||||
assert "organization_id" not in tags
|
||||
assert "organization_alias" not in tags
|
||||
assert tags["team_id"] == "team-1"
|
||||
Loading…
Reference in New Issue
Block a user