mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge defd5720d1 into 0eb6816081
This commit is contained in:
commit
9683178577
18 changed files with 638 additions and 61 deletions
|
|
@ -2327,6 +2327,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
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import re
|
||||
from collections.abc import ItemsView, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass, field
|
||||
|
|
@ -30,6 +31,9 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
_CAMEL_BOUNDARY: Final = re.compile(r"(?<=[a-z0-9])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])")
|
||||
|
||||
|
||||
def _object_items(data: Mapping[str, object]) -> ItemsView[str, object]:
|
||||
return data.items()
|
||||
|
||||
|
|
@ -44,6 +48,7 @@ class SensitiveDataMasker:
|
|||
mask_char: str = "*",
|
||||
mask_short_values: bool = True,
|
||||
extra_sensitive_patterns: AbstractSet[str] | None = None,
|
||||
split_camel_case: bool = False,
|
||||
):
|
||||
self.sensitive_patterns = (sensitive_patterns or _DEFAULT_SENSITIVE_PATTERNS) | (
|
||||
extra_sensitive_patterns or frozenset()
|
||||
|
|
@ -51,12 +56,15 @@ class SensitiveDataMasker:
|
|||
# If any key segment matches one of these, the key is not considered sensitive
|
||||
# even if it also matches a sensitive pattern. For example, "input_cost_per_token"
|
||||
# contains "token" but "cost" overrides that — it's a pricing field, not a secret.
|
||||
self.non_sensitive_overrides = non_sensitive_overrides or frozenset(("cost",))
|
||||
self.non_sensitive_overrides = (
|
||||
non_sensitive_overrides if non_sensitive_overrides is not None else frozenset(("cost",))
|
||||
)
|
||||
|
||||
self.visible_prefix = visible_prefix
|
||||
self.visible_suffix = visible_suffix
|
||||
self.mask_char = mask_char
|
||||
self.mask_short_values = mask_short_values
|
||||
self.split_camel_case = split_camel_case
|
||||
|
||||
def mask_value(self, value: str) -> str:
|
||||
value_str: Final = str(value)
|
||||
|
|
@ -82,11 +90,14 @@ class SensitiveDataMasker:
|
|||
if excluded_keys and key in excluded_keys:
|
||||
return False
|
||||
|
||||
key_lower: Final = str(key).lower()
|
||||
# Split on underscores/hyphens and check if any segment matches the pattern
|
||||
# This avoids false positives like "max_tokens" matching "token"
|
||||
# but still catches "api_key", "access_token", etc.
|
||||
key_segments: Final = key_lower.replace("-", "_").split("_")
|
||||
# Split on underscores and hyphens (and, when opted in, camelCase boundaries too) and check whether any
|
||||
# segment matches the pattern. This avoids false positives like "max_tokens" matching "token"
|
||||
# but still catches "api_key", "access_token" and, with split_camel_case, "apiKey", "APIKey" and "AWSSecret".
|
||||
key_lower: Final = str(key).lower().replace("-", "_")
|
||||
camel_lower: Final = (
|
||||
_CAMEL_BOUNDARY.sub("_", str(key)).lower().replace("-", "_") if self.split_camel_case else key_lower
|
||||
)
|
||||
key_segments: Final = frozenset(key_lower.split("_")) | frozenset(camel_lower.split("_"))
|
||||
|
||||
# If any segment matches a non-sensitive override, the key is not sensitive.
|
||||
# For example, "input_cost_per_token" contains "token" but also "cost",
|
||||
|
|
@ -257,10 +268,16 @@ def mask_sensitive_keys(data: Mapping[str, object], sensitive_fields: set[str])
|
|||
return masked
|
||||
|
||||
|
||||
def redact_credentials_in_payload(data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
def redact_credentials_in_payload(
|
||||
data: Mapping[str, object],
|
||||
masker: SensitiveDataMasker = _default_masker,
|
||||
*,
|
||||
dropped_keys: AbstractSet[str] = frozenset(),
|
||||
) -> Mapping[str, object]:
|
||||
"""Return a copy of ``data`` where every value under a credential-named key is
|
||||
replaced by the shared ``REDACTED`` marker, nested mappings are recursed into,
|
||||
and every other value is preserved by identity.
|
||||
replaced by the shared ``REDACTED`` marker, every entry whose key is in
|
||||
``dropped_keys`` is removed outright at any depth, nested mappings are recursed
|
||||
into, and every other value is preserved by identity.
|
||||
|
||||
Sensitive-key detection is delegated to the shared :class:`SensitiveDataMasker`,
|
||||
so the credential names stay in one place. Unlike
|
||||
|
|
@ -273,28 +290,31 @@ def redact_credentials_in_payload(data: Mapping[str, object]) -> Mapping[str, ob
|
|||
rather than passed through, so burying a credential deeper than the walk goes
|
||||
hides it instead of exposing it.
|
||||
"""
|
||||
return _redact_mapping(data, 0)
|
||||
return _PayloadRedactor(masker, frozenset(dropped_keys)).mapping(data, 0)
|
||||
|
||||
|
||||
def _redact_mapping(data: Mapping[str, object], depth: int) -> Mapping[str, object]:
|
||||
return {key: _redact_entry(key, value, depth) for key, value in data.items()}
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PayloadRedactor:
|
||||
masker: SensitiveDataMasker
|
||||
dropped_keys: frozenset[str]
|
||||
|
||||
def mapping(self, data: Mapping[str, object], depth: int) -> Mapping[str, object]:
|
||||
return {key: self.entry(key, value, depth) for key, value in data.items() if key not in self.dropped_keys}
|
||||
|
||||
def _redact_entry(key: str, value: object, depth: int) -> object:
|
||||
if value is not None and _default_masker.is_sensitive_key(key):
|
||||
return REDACTED
|
||||
if not isinstance(value, (Mapping, list, tuple)):
|
||||
return value
|
||||
if depth >= DEFAULT_MAX_RECURSE_DEPTH:
|
||||
return REDACTED
|
||||
if isinstance(value, Mapping):
|
||||
return _redact_mapping(value, depth + 1)
|
||||
return _redact_sequence(value, depth + 1)
|
||||
def entry(self, key: str, value: object, depth: int) -> object:
|
||||
if value is not None and self.masker.is_sensitive_key(key):
|
||||
return REDACTED
|
||||
if not isinstance(value, (Mapping, list, tuple)):
|
||||
return value
|
||||
if depth >= DEFAULT_MAX_RECURSE_DEPTH:
|
||||
return REDACTED
|
||||
if isinstance(value, Mapping):
|
||||
return self.mapping(value, depth + 1)
|
||||
return self.sequence(value, depth + 1)
|
||||
|
||||
|
||||
def _redact_sequence(values: Sequence[object], depth: int) -> Sequence[object]:
|
||||
redacted: Final = tuple(_redact_entry("", item, depth) for item in values)
|
||||
return redacted if isinstance(values, tuple) else list(redacted)
|
||||
def sequence(self, values: Sequence[object], depth: int) -> Sequence[object]:
|
||||
redacted: Final = tuple(self.entry("", item, depth) for item in values)
|
||||
return redacted if isinstance(values, tuple) else list(redacted)
|
||||
|
||||
|
||||
# Usage example:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -13,7 +13,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,
|
||||
|
|
@ -226,6 +226,8 @@ def _search_rows(keys: Sequence[str], response: SpendAnalyticsPaginatedResponse)
|
|||
|
||||
|
||||
def _csv_cell(value: object) -> object:
|
||||
if isinstance(value, Mapping):
|
||||
return json.dumps(value, separators=(",", ":"))
|
||||
if isinstance(value, str) and value.startswith(("=", "+", "-", "@", "\t", "\r")):
|
||||
return f"'{value}"
|
||||
return value
|
||||
|
|
@ -500,7 +502,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),
|
||||
|
|
|
|||
|
|
@ -110,9 +110,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
|
||||
|
||||
|
|
@ -891,9 +889,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(
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from dataclasses import replace
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
|
|
@ -11,8 +12,10 @@ from typing_extensions import assert_never
|
|||
from litellm import constants
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import UNKNOWN_MODEL_SPEND_LOG_MODEL
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker, redact_credentials_in_payload
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.daily_activity_sql import (
|
||||
EXPORT_METADATA_DROPPED_KEYS,
|
||||
ExportCursor,
|
||||
SqlQuery,
|
||||
adjust_dates_for_timezone,
|
||||
|
|
@ -140,6 +143,45 @@ def _daily_rows_table(
|
|||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
_EXPORT_UNSEPARATED_CREDENTIAL_NAMES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"apikey",
|
||||
"apitoken",
|
||||
"accesskey",
|
||||
"accesstoken",
|
||||
"authkey",
|
||||
"authtoken",
|
||||
"bearertoken",
|
||||
"clientsecret",
|
||||
"passwd",
|
||||
"privatekey",
|
||||
"pwd",
|
||||
"refreshtoken",
|
||||
"secretkey",
|
||||
"sessiontoken",
|
||||
"signingkey",
|
||||
"webhooksecret",
|
||||
}
|
||||
)
|
||||
_EXPORT_METADATA_MASKER: Final = SensitiveDataMasker(
|
||||
non_sensitive_overrides=frozenset(),
|
||||
extra_sensitive_patterns=_EXPORT_UNSEPARATED_CREDENTIAL_NAMES,
|
||||
split_camel_case=True,
|
||||
)
|
||||
|
||||
|
||||
def _redacted_metadata(metadata: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return redact_credentials_in_payload(metadata, _EXPORT_METADATA_MASKER, dropped_keys=EXPORT_METADATA_DROPPED_KEYS)
|
||||
|
||||
|
||||
def _with_redacted_metadata(row: ExportRow) -> ExportRow:
|
||||
if not row.key_metadata and not row.user_metadata:
|
||||
return row
|
||||
return replace(
|
||||
row, key_metadata=_redacted_metadata(row.key_metadata), user_metadata=_redacted_metadata(row.user_metadata)
|
||||
)
|
||||
|
||||
|
||||
def _next_export_cursor(batch: tuple[ExportRow, ...], export_type: ExportType) -> ExportCursor:
|
||||
last: Final = batch[-1]
|
||||
cursor_key: Final = (
|
||||
|
|
@ -258,15 +300,25 @@ 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
|
||||
yield _with_redacted_metadata(row)
|
||||
if len(batch) < batch_size:
|
||||
return
|
||||
cursor = _next_export_cursor(batch, export_type)
|
||||
|
|
|
|||
|
|
@ -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,12 +526,32 @@ def build_export_sql(
|
|||
)
|
||||
|
||||
|
||||
def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]:
|
||||
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}"
|
||||
EXPORT_METADATA_DROPPED_KEYS: Final[frozenset[str]] = CALLBACK_CONFIG_METADATA_SLOTS | SCIM_DIRECTORY_METADATA_KEYS
|
||||
_USER_METADATA_FIELD: Final = _metadata_object_field("u.metadata", EXPORT_METADATA_DROPPED_KEYS, "user_metadata")
|
||||
_KEY_METADATA_FIELD: Final = _metadata_object_field(
|
||||
"COALESCE(vt.metadata, dvt.metadata)", EXPORT_METADATA_DROPPED_KEYS, "key_metadata"
|
||||
)
|
||||
|
||||
|
||||
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 (
|
||||
"''",
|
||||
"NULL::text",
|
||||
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
|
||||
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email" + _EMPTY_METADATA_FIELDS,
|
||||
(),
|
||||
)
|
||||
if export_type is ExportType.DAILY_WITH_KEYS:
|
||||
|
|
@ -534,11 +560,12 @@ 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",
|
||||
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 (
|
||||
SELECT key_alias, user_id
|
||||
SELECT key_alias, user_id, metadata
|
||||
FROM "LiteLLM_DeletedVerificationToken"
|
||||
WHERE token = scoped.api_key
|
||||
ORDER BY deleted_at DESC
|
||||
|
|
@ -551,7 +578,7 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
|
|||
return (
|
||||
_MODEL_EXPR,
|
||||
"NULL::text",
|
||||
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
|
||||
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email" + _EMPTY_METADATA_FIELDS,
|
||||
(),
|
||||
)
|
||||
if export_type is ExportType.DAILY_WITH_USERS:
|
||||
|
|
@ -559,11 +586,11 @@ 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",
|
||||
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 (
|
||||
SELECT key_alias, user_id
|
||||
SELECT key_alias, user_id, metadata
|
||||
FROM "LiteLLM_DeletedVerificationToken"
|
||||
WHERE token = scoped.api_key
|
||||
ORDER BY deleted_at DESC
|
||||
|
|
|
|||
|
|
@ -19,6 +19,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"
|
||||
|
|
|
|||
|
|
@ -126,6 +126,8 @@ class ExportRow:
|
|||
failed_requests: int
|
||||
cache_read_input_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
key_metadata: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
|
||||
user_metadata: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -40,6 +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, empty_metadata_without_admin_view, redacts_credential_named_values], exercised_on: [chat_completions], source: "repositories/daily_activity_repository.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 with credential-named values replaced by REDACTED, 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"}
|
||||
|
|
|
|||
|
|
@ -59,6 +59,9 @@ class KeyMetadata(BaseModel):
|
|||
batch_enqueued_token_limit: int | None = None
|
||||
tag: str | None = None
|
||||
guardrails: list[str] | None = None
|
||||
webhook_token: str | None = None
|
||||
apiKey: str | None = None
|
||||
OpenAIAPIKey: str | None = None
|
||||
|
||||
|
||||
class ObjectPermission(BaseModel):
|
||||
|
|
@ -1683,6 +1686,7 @@ class UserNewBody(BaseModel):
|
|||
user_id: str | None = None
|
||||
auto_create_key: bool | None = None
|
||||
max_budget: float | None = None
|
||||
metadata: dict[str, str] | None = None
|
||||
|
||||
|
||||
class UserNewResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,203 @@
|
|||
"""The usage export carries the stored key and user metadata objects.
|
||||
|
||||
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 for
|
||||
callers with admin view, and returns {} for everyone else.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import provider_edge_base, unique_marker
|
||||
from e2e_metadata import Domain, Mode, Provider, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
KeyGenerateBody,
|
||||
KeyMetadata,
|
||||
LiteLLMParamsBody,
|
||||
UserNewBody,
|
||||
UserNewResponse,
|
||||
)
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from spend_e2e_client import SpendClient, unwrap
|
||||
from spend_reconciliation import BACKEND, INPUT_RATE, OUTPUT_RATE
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
EXPORT_ROUTE: Final = "/user/daily/activity/export"
|
||||
|
||||
|
||||
class UsageExportParams(BaseModel, frozen=True):
|
||||
start_date: str
|
||||
end_date: str
|
||||
api_key: str
|
||||
user_id: str | None = None
|
||||
export_type: str = "daily_with_keys"
|
||||
format: str = "json"
|
||||
|
||||
|
||||
class UsageExportKeyMetadata(BaseModel, frozen=True):
|
||||
tag: str | None = None
|
||||
webhook_token: str | None = None
|
||||
apiKey: str | None = None
|
||||
OpenAIAPIKey: str | None = None
|
||||
|
||||
|
||||
class UsageExportUserMetadata(BaseModel, frozen=True):
|
||||
department: str | None = None
|
||||
ldap_password: str | None = None
|
||||
clientSecret: str | None = None
|
||||
APIKEY: str | None = None
|
||||
|
||||
|
||||
class UsageExportRow(BaseModel, frozen=True):
|
||||
api_key: str | None = None
|
||||
user_id: str | None = None
|
||||
key_metadata: UsageExportKeyMetadata
|
||||
user_metadata: UsageExportUserMetadata
|
||||
|
||||
|
||||
class UsageExportRawMetadataRow(BaseModel, frozen=True):
|
||||
key_metadata: dict[str, object]
|
||||
user_metadata: dict[str, object]
|
||||
|
||||
|
||||
_EXPORT_ROWS: Final = TypeAdapter(list[UsageExportRow])
|
||||
_RAW_METADATA_ROWS: Final = TypeAdapter(list[UsageExportRawMetadataRow])
|
||||
|
||||
|
||||
class TestUsageExportMetadata:
|
||||
@pytest.mark.replayable
|
||||
@pytest.mark.covers("quota_management.spend_tracking.usage_export.carries_key_and_user_metadata")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(BACKEND,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_daily_with_keys_export_row_carries_key_and_user_metadata(
|
||||
self, client: SpendClient, resources: ResourceManager
|
||||
) -> None:
|
||||
started: Final = datetime.now(timezone.utc)
|
||||
marker: Final = unique_marker()
|
||||
base: Final = provider_edge_base("openai")
|
||||
model: Final = f"e2e-usage-export-{marker}"
|
||||
model_id: Final = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model=BACKEND,
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
api_base=None if base is None else f"{base}/v1",
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
user_id: Final = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=client.proxy.transport.master,
|
||||
json=UserNewBody(
|
||||
user_email=f"e2e-usage-export-{marker}@example.com",
|
||||
user_role="internal_user",
|
||||
user_id=f"e2e-usage-export-{marker}",
|
||||
metadata={
|
||||
"department": f"research-{marker}",
|
||||
"ldap_password": f"pw-{marker}",
|
||||
"clientSecret": f"cs-{marker}",
|
||||
"APIKEY": f"upper-{marker}",
|
||||
},
|
||||
),
|
||||
response_type=UserNewResponse,
|
||||
)
|
||||
).user_id
|
||||
resources.defer(lambda: client.delete_user(user_id))
|
||||
record: Final = client.generate_key_record(
|
||||
KeyGenerateBody(
|
||||
user_id=user_id,
|
||||
models=[model],
|
||||
key_alias=f"e2e-usage-export-{marker}",
|
||||
metadata=KeyMetadata(
|
||||
tag=f"cost-center-{marker}",
|
||||
webhook_token=f"whsec-{marker}",
|
||||
apiKey=f"ak-{marker}",
|
||||
OpenAIAPIKey=f"acr-{marker}",
|
||||
),
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_key(record.key))
|
||||
assert record.token, "/key/generate answered without the key's token hash"
|
||||
token: Final = record.token
|
||||
|
||||
response: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
record.key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=f"Reply with one word. {marker}")],
|
||||
max_completion_tokens=32,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert response.id, "successful response must have an ID"
|
||||
end: Final = datetime.now(timezone.utc) + timedelta(days=1)
|
||||
activity: Final = client.poll_daily_activity_for_key(
|
||||
token, start=started - timedelta(days=1), end=end, min_requests=1
|
||||
)
|
||||
assert activity is not None, f"/user/daily/activity never listed key {token} under api_keys"
|
||||
|
||||
exported: 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,
|
||||
),
|
||||
)
|
||||
assert exported.status_code == 200, f"{EXPORT_ROUTE} answered {exported.status_code}: {exported.body}"
|
||||
rows: Final = _EXPORT_ROWS.validate_json(exported.body)
|
||||
key_rows: Final = tuple(row for row in rows if row.api_key == token)
|
||||
assert key_rows, f"{EXPORT_ROUTE} returned no daily_with_keys row for {token}: {exported.body}"
|
||||
for row in key_rows:
|
||||
assert row.user_id == user_id, f"export row for {token} is owned by {row.user_id}, not {user_id}"
|
||||
assert row.key_metadata.tag == f"cost-center-{marker}", f"key metadata not exported: {row.key_metadata}"
|
||||
assert row.user_metadata.department == f"research-{marker}", (
|
||||
f"user metadata not exported: {row.user_metadata}"
|
||||
)
|
||||
masked = (
|
||||
row.key_metadata.webhook_token,
|
||||
row.key_metadata.apiKey,
|
||||
row.key_metadata.OpenAIAPIKey,
|
||||
row.user_metadata.ldap_password,
|
||||
row.user_metadata.clientSecret,
|
||||
row.user_metadata.APIKEY,
|
||||
)
|
||||
assert masked == ("REDACTED",) * 6, f"credential-named metadata values reached the export unmasked: {row}"
|
||||
|
||||
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 = _RAW_METADATA_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) == ({}, {}), (
|
||||
f"a caller without admin view got metadata objects: {row}"
|
||||
)
|
||||
|
|
@ -171,7 +171,17 @@ 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", "webhook_token": "whsec_1", '
|
||||
'"scim_roles": ["r"], "langfuse": {"apiKey": "lf", "logging": [{"callback_name": "langfuse"}]}, '
|
||||
'"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": []}', []),
|
||||
|
|
@ -198,8 +208,17 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str,
|
|||
),
|
||||
)
|
||||
cursor.execute(
|
||||
sql.SQL("INSERT INTO {} (user_id, user_email, models) VALUES (%s, %s, %s)").format(user_table),
|
||||
("user-1", "user@example.com", []),
|
||||
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", "ldap_password": "pw", "logging": [{"callback_name": "langfuse"}], '
|
||||
'"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",
|
||||
|
|
@ -236,8 +241,28 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
|
|||
"user-1",
|
||||
"user@example.com",
|
||||
)
|
||||
assert deleted_key_export.key_metadata == {"tags": ["archived"]}
|
||||
assert deleted_key_export.user_metadata == {"department": "research", "ldap_password": "REDACTED"}
|
||||
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"],
|
||||
"cost_center": "cc-1",
|
||||
"webhook_token": "REDACTED",
|
||||
"langfuse": {"apiKey": "REDACTED"},
|
||||
}
|
||||
assert active_key_export.user_metadata == {"department": "research", "ldap_password": "REDACTED"}
|
||||
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) == (
|
||||
|
|
@ -245,11 +270,17 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
|
|||
"user@example.com",
|
||||
273.0,
|
||||
)
|
||||
assert user_exports[0].user_metadata == {"department": "research", "ldap_password": "REDACTED"}
|
||||
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
|
||||
assert (daily_export[0].key_metadata, daily_export[0].user_metadata) == ({}, {})
|
||||
assert daily_export[0].flat_cost == 42.0
|
||||
|
||||
metadata: Final = await repository.key_metadata(frozenset(("key-a", "key-target")), None)
|
||||
|
|
@ -684,9 +715,7 @@ async def test_user_page_ranks_users_and_folds_null_ids() -> None:
|
|||
repository: Final = _repository(database)
|
||||
full_page: Final = await repository.user_page(scope, offset=0, limit=50)
|
||||
empty_page: Final = await repository.user_page(scope, offset=full_page.total_users + 1, limit=50)
|
||||
scoped_page: Final = await repository.user_page(
|
||||
replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50
|
||||
)
|
||||
scoped_page: Final = await repository.user_page(replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50)
|
||||
|
||||
assert tuple((row.user_id, row.spend) for row in full_page.rows) == (
|
||||
(None, 45.0),
|
||||
|
|
|
|||
|
|
@ -612,3 +612,68 @@ def test_redact_credentials_in_payload_leaves_a_realistic_tool_schema_intact():
|
|||
assert "REDACTED" not in str(result["tools"])
|
||||
assert result["tools"][0] == tool
|
||||
assert result["api_key"] == "REDACTED"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key",
|
||||
[
|
||||
"apiKey",
|
||||
"clientSecret",
|
||||
"accessToken",
|
||||
"secretAccessKey",
|
||||
"x-apiKey",
|
||||
"APIKey",
|
||||
"OpenAIAPIKey",
|
||||
"AWSSecret",
|
||||
"XAPIKey",
|
||||
],
|
||||
)
|
||||
def test_split_camel_case_opt_in_catches_camel_case_credential_names(key: str):
|
||||
assert SensitiveDataMasker(split_camel_case=True).is_sensitive_key(key) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["apiKey", "clientSecret", "tokenCount", "keyAlias", "accessLevel"])
|
||||
def test_default_masker_does_not_split_camel_case(key: str):
|
||||
assert SensitiveDataMasker().is_sensitive_key(key) is False
|
||||
|
||||
|
||||
def test_split_camel_case_keeps_pricing_and_count_fields_clear():
|
||||
masker: Final = SensitiveDataMasker(split_camel_case=True)
|
||||
|
||||
assert masker.is_sensitive_key("maxTokens") is False
|
||||
assert masker.is_sensitive_key("inputCostPerToken") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["sEt-CoOkIe", "AUTHORIZATION", "Api-Key", "X-API-KEY"])
|
||||
def test_split_camel_case_does_not_break_odd_cased_plain_names(key: str):
|
||||
assert SensitiveDataMasker(split_camel_case=True).is_sensitive_key(key) is True
|
||||
|
||||
|
||||
def test_empty_non_sensitive_overrides_disables_the_cost_exemption():
|
||||
strict: Final = SensitiveDataMasker(non_sensitive_overrides=frozenset())
|
||||
|
||||
assert SensitiveDataMasker().is_sensitive_key("cost_center_api_key") is False
|
||||
assert strict.is_sensitive_key("cost_center_api_key") is True
|
||||
assert strict.is_sensitive_key("cost_center") is False
|
||||
|
||||
|
||||
def test_redact_credentials_in_payload_drops_listed_keys_at_any_depth_with_the_given_masker():
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
|
||||
|
||||
payload: Final = {
|
||||
"logging": [{"callback_vars": {"langfuse_secret_key": "s"}}],
|
||||
"cost_center_api_key": "sk-cost",
|
||||
"nested": {"scim_roles": ["r"], "keep": 1, "deeper": [{"logging": {}, "apiKey": "k"}]},
|
||||
}
|
||||
|
||||
result: Final = redact_credentials_in_payload(
|
||||
payload,
|
||||
SensitiveDataMasker(non_sensitive_overrides=frozenset(), split_camel_case=True),
|
||||
dropped_keys=frozenset({"logging", "scim_roles"}),
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"cost_center_api_key": "REDACTED",
|
||||
"nested": {"keep": 1, "deeper": [{"apiKey": "REDACTED"}]},
|
||||
}
|
||||
assert redact_credentials_in_payload(payload)["cost_center_api_key"] == "sk-cost"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import csv
|
||||
import io
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, fields
|
||||
|
|
@ -266,6 +267,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(
|
||||
|
|
@ -375,7 +377,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):
|
||||
|
|
@ -397,6 +402,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} if include_metadata else {},
|
||||
user_metadata={"department": "research"} if include_metadata else {},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1407,6 +1414,9 @@ def test_export_routes_stream_csv_and_preserve_row_counts(
|
|||
assert len(records) == 7
|
||||
assert records[1][2] == "'=entity"
|
||||
assert records[1][4] == "'+key"
|
||||
columns: Final = dict(zip(records[0], records[1], strict=True))
|
||||
assert json.loads(columns["key_metadata"]) == {"cost_center": "=finance", "tier": 2}
|
||||
assert json.loads(columns["user_metadata"]) == {"department": "research"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("prefix", "query_name", "entity_id"), _ENTITY_CASES)
|
||||
|
|
@ -1427,6 +1437,41 @@ def test_export_routes_stream_json_arrays(
|
|||
records: Final = response.json()
|
||||
assert isinstance(records, list) and len(records) == 6, response.text
|
||||
assert records[0]["entity_alias"] == "=entity"
|
||||
assert records[0]["key_metadata"] == {"cost_center": "=finance", "tier": 2}
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -678,6 +678,61 @@ async def test_daily_rows_exclusion_without_entity_filter_keeps_null_entity_rows
|
|||
assert database.litellm_dailyteamspend.find_many_calls == [expected_where]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_export_redacts_credential_named_values_inside_metadata_objects() -> None:
|
||||
database: Final = _FakeDatabase(
|
||||
(
|
||||
(
|
||||
{
|
||||
**_export_row("key-1"),
|
||||
"key_metadata": {
|
||||
"cost_center": "cc-1",
|
||||
"cost_center_api_key": "sk-cost-1234567890",
|
||||
"api_key": "sk-live-1234567890",
|
||||
"webhook": {"url": "https://hooks.example", "auth_token": "t0k3n", "clientSecret": "cs"},
|
||||
"owners": [{"name": "ops", "password": "pw"}],
|
||||
"scim_roles": ["admin"],
|
||||
"langfuse": {"apiKey": "lf", "logging": [{"callback_vars": {"langfuse_secret_key": "s"}}]},
|
||||
"OpenAIAPIKey": "sk-acronym",
|
||||
"apikey": "sk-unseparated",
|
||||
"APIKEY": "sk-upper",
|
||||
"passwd": "pw-short",
|
||||
},
|
||||
"user_metadata": {
|
||||
"department": "research",
|
||||
"github_token": "ghp_secret",
|
||||
"logging": [{"callback_name": "langfuse"}],
|
||||
"profile": {"scim_enterprise": {"employeeNumber": "E-1"}, "team": "ops"},
|
||||
},
|
||||
},
|
||||
{**_export_row("key-2"), "key_metadata": {}, "user_metadata": {}},
|
||||
),
|
||||
)
|
||||
)
|
||||
repository, _ = _repository(database)
|
||||
|
||||
rows: Final = [row async for row in repository.export_rows(_scope(), export_type=ExportType.DAILY_WITH_KEYS)]
|
||||
|
||||
assert rows[0].key_metadata == {
|
||||
"cost_center": "cc-1",
|
||||
"cost_center_api_key": "REDACTED",
|
||||
"api_key": "REDACTED",
|
||||
"webhook": {"url": "https://hooks.example", "auth_token": "REDACTED", "clientSecret": "REDACTED"},
|
||||
"owners": [{"name": "ops", "password": "REDACTED"}],
|
||||
"langfuse": {"apiKey": "REDACTED"},
|
||||
"OpenAIAPIKey": "REDACTED",
|
||||
"apikey": "REDACTED",
|
||||
"APIKEY": "REDACTED",
|
||||
"passwd": "REDACTED",
|
||||
}
|
||||
assert rows[0].user_metadata == {
|
||||
"department": "research",
|
||||
"github_token": "REDACTED",
|
||||
"profile": {"team": "ops"},
|
||||
}
|
||||
assert (rows[1].key_metadata, rows[1].user_metadata) == ({}, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_export_is_lazy_and_uses_the_last_row_as_the_next_cursor(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 2)
|
||||
|
|
|
|||
|
|
@ -413,7 +413,7 @@ def test_query_builders_reject_nonpositive_limits(builder) -> None:
|
|||
def test_export_groups_by_requested_key_and_binds_cursor_after_scope(
|
||||
export_type: ExportType, group_key: str, key_filter: str, joins: str
|
||||
) -> None:
|
||||
query = build_export_sql(
|
||||
query: Final = build_export_sql(
|
||||
_scope(entity_ids=None),
|
||||
export_type=export_type,
|
||||
after=ExportCursor(date="2026-01-12", entity_id="user-2", group_key="group-3"),
|
||||
|
|
@ -444,16 +444,62 @@ def test_export_groups_by_requested_key_and_binds_cursor_after_scope(
|
|||
|
||||
@pytest.mark.parametrize("export_type", [ExportType.DAILY_WITH_KEYS, ExportType.DAILY_WITH_USERS])
|
||||
def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> None:
|
||||
query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
|
||||
query: Final = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
|
||||
|
||||
assert 'FROM "LiteLLM_DeletedVerificationToken"' in query.sql
|
||||
assert "ORDER BY deleted_at DESC" in query.sql
|
||||
assert "COALESCE(vt.user_id, dvt.user_id)" in query.sql
|
||||
|
||||
|
||||
_DROPPED_KEYS: Final = (
|
||||
"- 'callback_settings' - 'logging' - 'scim_enterprise' - 'scim_entitlements' - 'scim_roles' "
|
||||
"- 'secret_manager_settings'"
|
||||
)
|
||||
_REDACTED_KEY_METADATA: Final = (
|
||||
f"THEN (COALESCE(vt.metadata, dvt.metadata) {_DROPPED_KEYS})::text END), '{{}}')::jsonb AS key_metadata"
|
||||
)
|
||||
_REDACTED_USER_METADATA: Final = f"THEN (u.metadata {_DROPPED_KEYS})::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, _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_redacted_key_and_user_metadata_objects_for_admin_view(
|
||||
export_type: ExportType, key_metadata: str, user_metadata: str
|
||||
) -> None:
|
||||
query: Final = 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
|
||||
if export_type is ExportType.DAILY_WITH_KEYS:
|
||||
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: Final = 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(
|
||||
query: Final = build_export_sql(
|
||||
_scope(entity_ids=None),
|
||||
export_type=export_type,
|
||||
after=None,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue