mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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>
This commit is contained in:
parent
8c6625216b
commit
171fc49c28
4 changed files with 82 additions and 1 deletions
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -13,6 +13,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
|
||||
|
|
|
|||
60
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
60
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
"""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_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_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…
Add table
Reference in a new issue