mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(usage-export): only populate key and user metadata for admin view and strip credential and SCIM slots
This commit is contained in:
parent
a1b32c00ea
commit
e8e2fad99d
13 changed files with 190 additions and 53 deletions
|
|
@ -2312,6 +2312,9 @@ USAGE_MODEL_TOP_KEYS_MAX: Final[int] = 100
|
|||
USAGE_CACHE_LEAKAGE_KEYS_DEFAULT: Final[int] = 20
|
||||
USAGE_CACHE_LEAKAGE_KEYS_MAX: Final[int] = 100
|
||||
USAGE_EXPORT_BATCH_SIZE: Final[int] = 1000
|
||||
CALLBACK_CONFIG_METADATA_SLOTS: Final[frozenset[str]] = frozenset(
|
||||
{"logging", "callback_settings", "secret_manager_settings"}
|
||||
)
|
||||
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
|
||||
# declares no ptu_effective_from, bounding the scan for an open-ended window.
|
||||
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm import get_secret
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BAGGAGE_MAPPED_TAGS_METADATA_KEY,
|
||||
CALLBACK_CONFIG_METADATA_SLOTS,
|
||||
CLIENT_OUTPUT_CEILING_METADATA_KEY,
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY,
|
||||
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH,
|
||||
|
|
@ -52,7 +53,7 @@ _CALLBACK_VAR_ENCRYPTED_PREFIX: Final = CALLBACK_VAR_ENCRYPTED_PREFIX
|
|||
# Metadata slots that hold operator-configured callback and secret-manager setup
|
||||
# (and therefore integration credentials). Resolved from UserAPIKeyAuth during
|
||||
# pre-call setup, never read back off the copies stamped into request metadata.
|
||||
_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings", "secret_manager_settings"})
|
||||
_CALLBACK_CONFIG_SLOTS: Final = CALLBACK_CONFIG_METADATA_SLOTS
|
||||
|
||||
blue_color_code: Final = "\033[94m"
|
||||
reset_color_code: Final = "\033[0m"
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from fastapi.responses import StreamingResponse
|
|||
|
||||
from litellm import constants
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
InvalidDateRange,
|
||||
|
|
@ -497,7 +497,11 @@ def _register_export_route(router: APIRouter, resolver: EntityScopeResolver, pre
|
|||
prisma_client,
|
||||
user_aggregated=False,
|
||||
)
|
||||
rows: Final = repository.export_rows(resolved.scope, export_type=export_type)
|
||||
rows: Final = repository.export_rows(
|
||||
resolved.scope,
|
||||
export_type=export_type,
|
||||
include_metadata=user_api_key_has_admin_view(user_api_key_dict),
|
||||
)
|
||||
first_row: Final = await anext(rows, None)
|
||||
return StreamingResponse(
|
||||
_stream_export_rows(first_row, rows, file_format),
|
||||
|
|
|
|||
|
|
@ -106,9 +106,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import (
|
|||
UserUpdateResult,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.scim_v2 import (
|
||||
SCIM_ENTERPRISE_METADATA_KEY,
|
||||
SCIM_ENTITLEMENTS_METADATA_KEY,
|
||||
SCIM_ROLES_METADATA_KEY,
|
||||
SCIM_DIRECTORY_METADATA_KEYS,
|
||||
)
|
||||
from litellm.types.utils import BudgetConfig
|
||||
|
||||
|
|
@ -882,9 +880,7 @@ async def _get_user_info_teams(
|
|||
return team_list, teams_1
|
||||
|
||||
|
||||
_SCIM_DIRECTORY_METADATA_KEYS: Final = frozenset(
|
||||
{SCIM_ENTERPRISE_METADATA_KEY, SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY}
|
||||
)
|
||||
_SCIM_DIRECTORY_METADATA_KEYS: Final = SCIM_DIRECTORY_METADATA_KEYS
|
||||
|
||||
|
||||
def _redact_scim_enterprise_metadata(
|
||||
|
|
|
|||
|
|
@ -258,12 +258,22 @@ class DailyActivityRepository:
|
|||
query: Final = build_cache_leakage_keys_sql(scope, limit=limit)
|
||||
return _KEY_SPEND_ADAPTER.validate_python(await self._query(query))
|
||||
|
||||
async def export_rows(self, scope: DailyActivityScope, *, export_type: ExportType) -> AsyncIterator[ExportRow]:
|
||||
async def export_rows(
|
||||
self, scope: DailyActivityScope, *, export_type: ExportType, include_metadata: bool = False
|
||||
) -> AsyncIterator[ExportRow]:
|
||||
batch_size: Final = constants.USAGE_EXPORT_BATCH_SIZE
|
||||
cursor: ExportCursor | None = None # rebind-ok: each page advances the export keyset cursor
|
||||
while True:
|
||||
batch: tuple[ExportRow, ...] = _EXPORT_ADAPTER.validate_python(
|
||||
await self._query(build_export_sql(scope, export_type=export_type, after=cursor, batch_size=batch_size))
|
||||
await self._query(
|
||||
build_export_sql(
|
||||
scope,
|
||||
export_type=export_type,
|
||||
after=cursor,
|
||||
batch_size=batch_size,
|
||||
include_metadata=include_metadata,
|
||||
)
|
||||
)
|
||||
)
|
||||
for row in batch:
|
||||
yield row
|
||||
|
|
|
|||
|
|
@ -8,8 +8,9 @@ from typing import Final
|
|||
from typing_extensions import assert_never
|
||||
|
||||
from litellm import constants
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.constants import CALLBACK_CONFIG_METADATA_SLOTS, PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.repositories.daily_activity_buckets import UNATTRIBUTED_API_KEY, UNKNOWN_MODEL_GROUP
|
||||
from litellm.types.proxy.management_endpoints.scim_v2 import SCIM_DIRECTORY_METADATA_KEYS
|
||||
from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType
|
||||
|
||||
_API_KEY_ROLLED_UP_BIT: Final = 32
|
||||
|
|
@ -447,11 +448,16 @@ def build_cache_leakage_keys_sql(scope: DailyActivityScope, *, limit: int) -> Sq
|
|||
|
||||
|
||||
def build_export_sql(
|
||||
scope: DailyActivityScope, *, export_type: ExportType, after: ExportCursor | None, batch_size: int
|
||||
scope: DailyActivityScope,
|
||||
*,
|
||||
export_type: ExportType,
|
||||
after: ExportCursor | None,
|
||||
batch_size: int,
|
||||
include_metadata: bool = False,
|
||||
) -> SqlQuery:
|
||||
_bounded_limit(batch_size)
|
||||
where_clause, where_params = build_where_clause(scope)
|
||||
group_key, output_key, user_fields, type_joins = _export_grouping(export_type)
|
||||
group_key, output_key, user_fields, type_joins = _export_grouping(export_type, include_metadata=include_metadata)
|
||||
grouping_keys: Final = (
|
||||
f"scoped.date, COALESCE(scoped.\"{scope.entity_id_field}\", '')",
|
||||
*((group_key,) if export_type is not ExportType.DAILY else ()),
|
||||
|
|
@ -520,18 +526,26 @@ def build_export_sql(
|
|||
)
|
||||
|
||||
|
||||
_EMPTY_METADATA_FIELDS: Final = ", '{}'::jsonb AS key_metadata, '{}'::jsonb AS user_metadata"
|
||||
_USER_METADATA_FIELD: Final = (
|
||||
"COALESCE(MIN(CASE WHEN jsonb_typeof(u.metadata) = 'object' THEN u.metadata::text END), '{}')::jsonb "
|
||||
"AS user_metadata"
|
||||
)
|
||||
_KEY_METADATA_FIELD: Final = (
|
||||
"COALESCE(MIN(CASE WHEN jsonb_typeof(COALESCE(vt.metadata, dvt.metadata)) = 'object' "
|
||||
"THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata"
|
||||
def _metadata_object_field(expression: str, redacted_keys: frozenset[str], alias: str) -> str:
|
||||
redacted: Final = " - ".join((expression, *(f"'{key}'" for key in sorted(redacted_keys))))
|
||||
return (
|
||||
f"COALESCE(MIN(CASE WHEN jsonb_typeof({expression}) = 'object' THEN ({redacted})::text END), '{{}}')::jsonb "
|
||||
f"AS {alias}"
|
||||
)
|
||||
|
||||
|
||||
_EMPTY_KEY_METADATA_FIELD: Final = "'{}'::jsonb AS key_metadata"
|
||||
_EMPTY_USER_METADATA_FIELD: Final = "'{}'::jsonb AS user_metadata"
|
||||
_EMPTY_METADATA_FIELDS: Final = f", {_EMPTY_KEY_METADATA_FIELD}, {_EMPTY_USER_METADATA_FIELD}"
|
||||
_USER_METADATA_FIELD: Final = _metadata_object_field("u.metadata", SCIM_DIRECTORY_METADATA_KEYS, "user_metadata")
|
||||
_KEY_METADATA_FIELD: Final = _metadata_object_field(
|
||||
"COALESCE(vt.metadata, dvt.metadata)", CALLBACK_CONFIG_METADATA_SLOTS, "key_metadata"
|
||||
)
|
||||
|
||||
|
||||
def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]:
|
||||
def _export_grouping(export_type: ExportType, *, include_metadata: bool) -> tuple[str, str, str, tuple[str, ...]]:
|
||||
key_metadata: Final = _KEY_METADATA_FIELD if include_metadata else _EMPTY_KEY_METADATA_FIELD
|
||||
user_metadata: Final = _USER_METADATA_FIELD if include_metadata else _EMPTY_USER_METADATA_FIELD
|
||||
if export_type is ExportType.DAILY:
|
||||
return (
|
||||
"''",
|
||||
|
|
@ -545,10 +559,8 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
|
|||
_API_KEY_EXPR,
|
||||
"MAX(CASE WHEN scoped.api_key = '' THEN 'Unattributed' "
|
||||
"ELSE COALESCE(vt.key_alias, dvt.key_alias) END) AS key_alias, "
|
||||
"MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email, "
|
||||
+ _KEY_METADATA_FIELD
|
||||
+ ", "
|
||||
+ _USER_METADATA_FIELD,
|
||||
f"MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email, "
|
||||
f"{key_metadata}, {user_metadata}",
|
||||
(
|
||||
"LEFT JOIN \"LiteLLM_VerificationToken\" vt ON vt.token = scoped.api_key AND scoped.api_key <> ''",
|
||||
"""LEFT JOIN LATERAL (
|
||||
|
|
@ -573,7 +585,7 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
|
|||
"COALESCE(vt.user_id, dvt.user_id, '')",
|
||||
"NULL::text",
|
||||
"NULL::text AS key_alias, MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, "
|
||||
"MAX(u.user_email) AS user_email, '{}'::jsonb AS key_metadata, " + _USER_METADATA_FIELD,
|
||||
f"MAX(u.user_email) AS user_email, {_EMPTY_KEY_METADATA_FIELD}, {user_metadata}",
|
||||
(
|
||||
'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key',
|
||||
"""LEFT JOIN LATERAL (
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:ent
|
|||
SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise"
|
||||
SCIM_ENTITLEMENTS_METADATA_KEY: Final = "scim_entitlements"
|
||||
SCIM_ROLES_METADATA_KEY: Final = "scim_roles"
|
||||
SCIM_DIRECTORY_METADATA_KEYS: Final[frozenset[str]] = frozenset(
|
||||
{SCIM_ENTERPRISE_METADATA_KEY, SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY}
|
||||
)
|
||||
|
||||
SCIM_MANAGED_TEAM_METADATA_KEY: Final = "scim_managed"
|
||||
SCIM_TEAM_DATA_METADATA_KEY: Final = "scim_data"
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@
|
|||
- {id: quota_management.budget.fallback.routes_to_fallback, module: quota_management, tier: P1, behavior: budget, variant: fallback, assertions: [routes_to_fallback], exercised_on: [messages], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "budget_fallbacks reroute to the fallback model once the primary's budget is exhausted"}
|
||||
- {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"}
|
||||
- {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"}
|
||||
- {id: quota_management.spend_tracking.usage_export.carries_key_and_user_metadata, module: quota_management, tier: P1, behavior: spend_tracking, variant: usage_export, assertions: [carries_key_and_user_metadata], exercised_on: [chat_completions], source: "repositories/daily_activity_sql.py / LIT-8524", rationale: "The daily_with_keys usage export row carries the key's and the owning user's stored metadata objects"}
|
||||
- {id: quota_management.spend_tracking.usage_export.carries_key_and_user_metadata, module: quota_management, tier: P1, behavior: spend_tracking, variant: usage_export, assertions: [carries_key_and_user_metadata, empty_metadata_without_admin_view], exercised_on: [chat_completions], source: "repositories/daily_activity_sql.py / LIT-8524", rationale: "The daily_with_keys usage export row carries the key's and the owning user's stored metadata objects for callers with admin view and {} for everyone else"}
|
||||
- {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"}
|
||||
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
|
||||
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ A customer charging usage back by cost center keeps that cost center in the
|
|||
virtual key's metadata and the owner's user metadata. Before LIT-8524 the
|
||||
Usage page export only carried the key alias, user id and email, so every
|
||||
exported row needed a /key/info and /user/info round trip to recover the
|
||||
metadata. The export now joins both objects onto each daily_with_keys row.
|
||||
metadata. The export now joins both objects onto each daily_with_keys row for
|
||||
callers with admin view, and returns {} for everyone else.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -38,6 +39,7 @@ class UsageExportParams(BaseModel):
|
|||
start_date: str
|
||||
end_date: str
|
||||
api_key: str
|
||||
user_id: str | None = None
|
||||
export_type: str = "daily_with_keys"
|
||||
format: str = "json"
|
||||
|
||||
|
|
@ -150,3 +152,21 @@ class TestUsageExportMetadata:
|
|||
assert row.user_metadata.department == f"research-{marker}", (
|
||||
f"user metadata not exported: {row.user_metadata}"
|
||||
)
|
||||
|
||||
own_view: Final = client.proxy.transport.probe(
|
||||
EXPORT_ROUTE,
|
||||
params=UsageExportParams(
|
||||
start_date=(started - timedelta(days=1)).strftime("%Y-%m-%d"),
|
||||
end_date=end.strftime("%Y-%m-%d"),
|
||||
api_key=token,
|
||||
user_id=user_id,
|
||||
),
|
||||
headers=client.proxy.transport.bearer(record.key),
|
||||
)
|
||||
assert own_view.status_code == 200, f"{EXPORT_ROUTE} answered {own_view.status_code}: {own_view.body}"
|
||||
own_rows: Final = _EXPORT_ROWS.validate_json(own_view.body)
|
||||
assert own_rows, f"{EXPORT_ROUTE} returned no rows for the key's own user: {own_view.body}"
|
||||
for row in own_rows:
|
||||
assert (row.key_metadata, row.user_metadata) == (UsageExportKeyMetadata(), UsageExportUserMetadata()), (
|
||||
f"a caller without admin view got metadata objects: {row}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -171,7 +171,16 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str,
|
|||
"INSERT INTO {} (token, key_alias, team_id, user_id, metadata, models) VALUES (%s, %s, %s, %s, %s, %s)"
|
||||
).format(verification_token_table),
|
||||
(
|
||||
("key-a", "alias-a", "team-1", "user-1", '{"tags": ["blue", "gold"]}', []),
|
||||
(
|
||||
"key-a",
|
||||
"alias-a",
|
||||
"team-1",
|
||||
"user-1",
|
||||
'{"tags": ["blue", "gold"], "cost_center": "cc-1", '
|
||||
'"logging": [{"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk"}}], '
|
||||
'"callback_settings": {"otel": {"headers": "x"}}, "secret_manager_settings": {"vault": "v"}}',
|
||||
[],
|
||||
),
|
||||
("key-b", "alias-b", "team-1", "user-1", '{"tags": []}', []),
|
||||
("key-c", "alias-c", "team-1", "user-1", '{"tags": []}', []),
|
||||
("key-cache", "alias-cache", "team-1", "user-1", '{"tags": []}', []),
|
||||
|
|
@ -201,7 +210,13 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str,
|
|||
sql.SQL("INSERT INTO {} (user_id, user_email, models, metadata) VALUES (%s, %s, %s, %s)").format(
|
||||
user_table
|
||||
),
|
||||
("user-1", "user@example.com", [], '{"department": "research"}'),
|
||||
(
|
||||
"user-1",
|
||||
"user@example.com",
|
||||
[],
|
||||
'{"department": "research", "scim_enterprise": {"employeeNumber": "E-1"}, '
|
||||
'"scim_entitlements": ["ent"], "scim_roles": ["role"]}',
|
||||
),
|
||||
)
|
||||
cursor.execute(
|
||||
sql.SQL("INSERT INTO {} (team_id, team_alias, admins, members, models) VALUES (%s, %s, %s, %s, %s)").format(
|
||||
|
|
|
|||
|
|
@ -219,7 +219,12 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
|
|||
assert tuple(row.api_key for row in leakage_keys) == ("key-cache", "key-c")
|
||||
|
||||
exports: Final = tuple(
|
||||
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS)]
|
||||
[
|
||||
row
|
||||
async for row in repository.export_rows(
|
||||
team_scope, export_type=ExportType.DAILY_WITH_KEYS, include_metadata=True
|
||||
)
|
||||
]
|
||||
)
|
||||
assert tuple(row.api_key for row in exports) == (
|
||||
"key-a",
|
||||
|
|
@ -239,10 +244,20 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
|
|||
assert deleted_key_export.key_metadata == {"tags": ["archived"]}
|
||||
assert deleted_key_export.user_metadata == {"department": "research"}
|
||||
active_key_export: Final = next(row for row in exports if row.api_key == "key-a")
|
||||
assert active_key_export.key_metadata == {"tags": ["blue", "gold"]}
|
||||
assert active_key_export.key_metadata == {"tags": ["blue", "gold"], "cost_center": "cc-1"}
|
||||
assert active_key_export.user_metadata == {"department": "research"}
|
||||
unprivileged_exports: Final = tuple(
|
||||
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS)]
|
||||
)
|
||||
assert tuple(row.api_key for row in unprivileged_exports) == tuple(row.api_key for row in exports)
|
||||
assert all((row.key_metadata, row.user_metadata) == ({}, {}) for row in unprivileged_exports)
|
||||
user_exports: Final = tuple(
|
||||
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_USERS)]
|
||||
[
|
||||
row
|
||||
async for row in repository.export_rows(
|
||||
team_scope, export_type=ExportType.DAILY_WITH_USERS, include_metadata=True
|
||||
)
|
||||
]
|
||||
)
|
||||
assert len(user_exports) == 1
|
||||
assert (user_exports[0].user_id, user_exports[0].user_email, user_exports[0].spend) == (
|
||||
|
|
@ -253,7 +268,10 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
|
|||
assert user_exports[0].user_metadata == {"department": "research"}
|
||||
assert user_exports[0].key_metadata == {}
|
||||
daily_export: Final = tuple(
|
||||
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY)]
|
||||
[
|
||||
row
|
||||
async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY, include_metadata=True)
|
||||
]
|
||||
)
|
||||
assert len(daily_export) == 1
|
||||
assert daily_export[0].spend == 1273.0
|
||||
|
|
|
|||
|
|
@ -260,6 +260,7 @@ class _FakeRepository:
|
|||
self.user_page_call: tuple[DailyActivityScope, int, int] | None = None
|
||||
self.user_metadata = AsyncMock(side_effect=self._user_metadata)
|
||||
self.export_rows_error: Exception | None = None
|
||||
self.export_rows_include_metadata: bool | None = None
|
||||
|
||||
def _matching_rows(self, scope: DailyActivityScope) -> tuple[_Activity, ...]:
|
||||
return tuple(
|
||||
|
|
@ -369,7 +370,10 @@ class _FakeRepository:
|
|||
for key in api_keys
|
||||
}
|
||||
|
||||
async def export_rows(self, scope: DailyActivityScope, *, export_type: ExportType) -> AsyncIterator[ExportRow]:
|
||||
async def export_rows(
|
||||
self, scope: DailyActivityScope, *, export_type: ExportType, include_metadata: bool = False
|
||||
) -> AsyncIterator[ExportRow]:
|
||||
self.export_rows_include_metadata = include_metadata
|
||||
if self.export_rows_error is not None:
|
||||
raise self.export_rows_error
|
||||
for row in self._matching_rows(scope):
|
||||
|
|
@ -391,8 +395,8 @@ class _FakeRepository:
|
|||
failed_requests=row.failed_requests,
|
||||
cache_read_input_tokens=row.cache_read_input_tokens,
|
||||
cache_creation_input_tokens=row.cache_creation_input_tokens,
|
||||
key_metadata={"cost_center": "=finance", "tier": 2},
|
||||
user_metadata={"department": "research"},
|
||||
key_metadata={"cost_center": "=finance", "tier": 2} if include_metadata else {},
|
||||
user_metadata={"department": "research"} if include_metadata else {},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1190,6 +1194,39 @@ def test_export_routes_stream_json_arrays(
|
|||
assert records[0]["user_metadata"] == {"department": "research"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("role", "expects_metadata"),
|
||||
[
|
||||
(LitellmUserRoles.PROXY_ADMIN, True),
|
||||
(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True),
|
||||
(LitellmUserRoles.INTERNAL_USER, False),
|
||||
],
|
||||
)
|
||||
def test_export_metadata_objects_only_for_callers_with_admin_view(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
role: LitellmUserRoles,
|
||||
expects_metadata: bool,
|
||||
) -> None:
|
||||
client, repository = daily_activity_client
|
||||
response: Final = client.get(
|
||||
"/user/daily/activity/export",
|
||||
params={
|
||||
**_entity_params("user_id", "user-a"),
|
||||
"export_type": ExportType.DAILY_WITH_KEYS.value,
|
||||
"format": "json",
|
||||
},
|
||||
headers={"x-user-role": role.value, "x-user-id": "user-a"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert repository.export_rows_include_metadata is expects_metadata
|
||||
records: Final = response.json()
|
||||
assert records, response.text
|
||||
expected_key: Final = {"cost_center": "=finance", "tier": 2} if expects_metadata else {}
|
||||
expected_user: Final = {"department": "research"} if expects_metadata else {}
|
||||
assert all(record["key_metadata"] == expected_key for record in records), response.text
|
||||
assert all(record["user_metadata"] == expected_user for record in records), response.text
|
||||
|
||||
|
||||
def test_export_first_row_error_returns_json_error_before_streaming(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -451,27 +451,33 @@ def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> Non
|
|||
assert "COALESCE(vt.user_id, dvt.user_id)" in query.sql
|
||||
|
||||
|
||||
_REDACTED_KEY_METADATA: Final = (
|
||||
"THEN (COALESCE(vt.metadata, dvt.metadata) - 'callback_settings' - 'logging' - 'secret_manager_settings')"
|
||||
"::text END), '{}')::jsonb AS key_metadata"
|
||||
)
|
||||
_REDACTED_USER_METADATA: Final = (
|
||||
"THEN (u.metadata - 'scim_enterprise' - 'scim_entitlements' - 'scim_roles')::text END), '{}')::jsonb "
|
||||
"AS user_metadata"
|
||||
)
|
||||
_EMPTY_KEY_METADATA: Final = "'{}'::jsonb AS key_metadata"
|
||||
_EMPTY_USER_METADATA: Final = "'{}'::jsonb AS user_metadata"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("export_type", "key_metadata", "user_metadata"),
|
||||
[
|
||||
(
|
||||
ExportType.DAILY_WITH_KEYS,
|
||||
"THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata",
|
||||
"THEN u.metadata::text END), '{}')::jsonb AS user_metadata",
|
||||
),
|
||||
(
|
||||
ExportType.DAILY_WITH_USERS,
|
||||
"'{}'::jsonb AS key_metadata",
|
||||
"THEN u.metadata::text END), '{}')::jsonb AS user_metadata",
|
||||
),
|
||||
(ExportType.DAILY, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"),
|
||||
(ExportType.DAILY_WITH_MODELS, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"),
|
||||
(ExportType.DAILY_WITH_KEYS, _REDACTED_KEY_METADATA, _REDACTED_USER_METADATA),
|
||||
(ExportType.DAILY_WITH_USERS, _EMPTY_KEY_METADATA, _REDACTED_USER_METADATA),
|
||||
(ExportType.DAILY, _EMPTY_KEY_METADATA, _EMPTY_USER_METADATA),
|
||||
(ExportType.DAILY_WITH_MODELS, _EMPTY_KEY_METADATA, _EMPTY_USER_METADATA),
|
||||
],
|
||||
)
|
||||
def test_export_selects_stored_key_and_user_metadata_objects(
|
||||
def test_export_selects_redacted_key_and_user_metadata_objects_for_admin_view(
|
||||
export_type: ExportType, key_metadata: str, user_metadata: str
|
||||
) -> None:
|
||||
query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
|
||||
query = build_export_sql(
|
||||
_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2, include_metadata=True
|
||||
)
|
||||
|
||||
assert key_metadata in query.sql
|
||||
assert user_metadata in query.sql
|
||||
|
|
@ -479,6 +485,18 @@ def test_export_selects_stored_key_and_user_metadata_objects(
|
|||
assert "SELECT key_alias, user_id, metadata" in query.sql
|
||||
|
||||
|
||||
@pytest.mark.parametrize("export_type", tuple(ExportType))
|
||||
def test_export_without_metadata_selects_empty_objects_and_never_reads_metadata_columns(
|
||||
export_type: ExportType,
|
||||
) -> None:
|
||||
query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
|
||||
|
||||
assert _EMPTY_KEY_METADATA in query.sql
|
||||
assert _EMPTY_USER_METADATA in query.sql
|
||||
assert "u.metadata" not in query.sql
|
||||
assert "vt.metadata" not in query.sql
|
||||
|
||||
|
||||
@pytest.mark.parametrize("export_type", tuple(ExportType))
|
||||
def test_export_without_cursor_omits_cursor_predicate_and_parameters(export_type: ExportType) -> None:
|
||||
query = build_export_sql(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue