mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy): let a key's logging entry reference a proxy credential (#50572)
* feat(proxy): let a key's logging entry reference a proxy credential
A virtual key's metadata.logging entry can name a proxy credential with
litellm_credential_name. At request time the credential's values, which config
credential_list resolves through the secret manager at boot, are merged under the
entry's callback_vars, with the entry's own values winning. An entry whose
credential is not defined is dropped, so it never falls back to the proxy's own
callback environment.
Only proxy admins may change the logging entries of a key when the stored or
requested list holds a reference (403), and a reference to an undefined
credential is a 400. The check runs on key generate, update, team bulk update,
regenerate and /user/new, and compares decrypted entries so re-saving metadata
read back from /key/info is not a change.
Resolves LIT-10061
* fix(proxy): drop a key logging entry whose referenced credential holds no logging values
* fix(proxy): only require newly added credential references to exist on key writes
* fix(proxy): log a dropped credential reference as an error so /key/health reports it
* fix(proxy): gate credential references on bulk-created user keys
* fix(proxy): keep failing requests on keys with non-object logging metadata
A key whose stored metadata.logging is not a list of objects failed every
request with a 500 before this branch. The credential reference rewrite
turned those values into no entries, so the request quietly succeeded
without the key's logging. Raise a TypeError for those shapes again so the
request path behaves as it did on main, while empty values and objects with
unknown fields still go through the existing skip.
* test(proxy): audit cells for credential references across surfaces, v1, dropped entries and malformed logging
The malformed logging cell expects the base behavior that 7ac4ab1356
restored, a 500 for a key whose stored metadata.logging is not a list of
objects and a 200 for an empty list
* fix(proxy): keep a 500 for teams with non-object logging metadata
On main, a team whose stored metadata.logging held non-object items failed
chat requests with a 500, because AddTeamCallback(**item) raised a TypeError.
This branch switched to AddTeamCallback.model_validate, whose before-validator
raises AttributeError on those items, and the proxy maps AttributeError to a
400 invalid_request_error. Raise a TypeError for non-mapping items before
validation so the status matches main again, as the key path already does.
This commit is contained in:
jesus-berri
GitHub
parent
e5ea9502ba
commit
c76a95fc7e
14 changed files with 929 additions and 24 deletions
|
|
@ -2101,6 +2101,11 @@ class BlockModelRequest(LiteLLMPydanticObjectBase):
|
|||
model_id: str # required
|
||||
|
||||
|
||||
CALLBACK_VAR_NAMES: Final = frozenset(
|
||||
{*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR}
|
||||
)
|
||||
|
||||
|
||||
class AddTeamCallback(LiteLLMPydanticObjectBase):
|
||||
callback_name: str
|
||||
callback_type: Literal["success", "failure", "success_and_failure"] | None = "success_and_failure"
|
||||
|
|
@ -2110,14 +2115,9 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
|
|||
@classmethod
|
||||
def validate_callback_vars(cls, values):
|
||||
callback_vars: Final = values.get("callback_vars", {})
|
||||
valid_keys: Final = {
|
||||
*StandardCallbackDynamicParams.__annotations__,
|
||||
CAPTURE_MESSAGE_CONTENT_VAR,
|
||||
EXCLUDED_SERVICES_VAR,
|
||||
}
|
||||
for key, value in callback_vars.items():
|
||||
if key not in valid_keys:
|
||||
raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}")
|
||||
if key not in CALLBACK_VAR_NAMES:
|
||||
raise ValueError(f"Invalid callback variable: {key}. Must be one of {set(CALLBACK_VAR_NAMES)}")
|
||||
if key in (CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR) and not isinstance(value, str):
|
||||
raise ValueError(f"{key} must be a string, got {type(value).__name__}")
|
||||
callback_vars[key] = str(value)
|
||||
|
|
|
|||
|
|
@ -7,9 +7,13 @@ per-integration checks here.
|
|||
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
|
||||
from typing import Final, Literal, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR
|
||||
|
||||
_NEWRELIC_CALLBACK: Final = "newrelic"
|
||||
|
|
@ -276,7 +280,7 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
|
|||
"""Validate every ``logging`` entry of a team/key metadata payload."""
|
||||
if not metadata:
|
||||
return None
|
||||
entries: Final = _object_sequence(metadata.get("logging"))
|
||||
entries: Final = object_sequence(metadata.get("logging"))
|
||||
if entries is None:
|
||||
return None
|
||||
entry_error: Final = next(
|
||||
|
|
@ -302,11 +306,11 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
|
|||
|
||||
def stored_callback_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]:
|
||||
"""The callback name and vars of every stored ``logging`` entry."""
|
||||
entries: Final = _object_sequence(logging_entries) or ()
|
||||
entries: Final = object_sequence(logging_entries) or ()
|
||||
return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries)
|
||||
|
||||
|
||||
def _object_sequence(value: object) -> Sequence[object] | None:
|
||||
def object_sequence(value: object) -> Sequence[object] | None:
|
||||
return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None
|
||||
|
||||
|
||||
|
|
@ -445,3 +449,72 @@ def _newrelic_config_error(callback_vars: Mapping[str, str]) -> str | None:
|
|||
if region is not None and (not api_key or api_key == "None"):
|
||||
return "newrelic_region requires newrelic_api_key; the region rides the team's own key."
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CredentialReferenceRejection:
|
||||
status_code: Literal[400, 403]
|
||||
message: str
|
||||
|
||||
|
||||
def credential_is_defined(credential_name: str) -> bool:
|
||||
return CredentialAccessor.find_credential(credential_name) is not None
|
||||
|
||||
|
||||
_STRING_KEYED: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def string_keyed(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _STRING_KEYED.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _logging_entries(metadata: object) -> object:
|
||||
return string_keyed(metadata).get("logging")
|
||||
|
||||
|
||||
def _entry_credential_name(entry: object) -> str | None:
|
||||
match entry:
|
||||
case {"litellm_credential_name": str() as credential_name}:
|
||||
return credential_name
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def logging_credential_names(logging_entries: object) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
name for entry in object_sequence(logging_entries) or () if (name := _entry_credential_name(entry)) is not None
|
||||
)
|
||||
|
||||
|
||||
def credential_reference_rejection(
|
||||
requested_metadata: object,
|
||||
stored_logging: object,
|
||||
*,
|
||||
is_proxy_admin: bool,
|
||||
is_defined: Callable[[str], bool] = credential_is_defined,
|
||||
) -> CredentialReferenceRejection | None:
|
||||
"""Gate a key metadata write whose logging entries reference a proxy credential.
|
||||
|
||||
``requested_metadata`` is None when the write leaves the key's metadata untouched. The admin check runs before
|
||||
the existence check so a non-admin cannot probe which credential names are defined.
|
||||
"""
|
||||
if requested_metadata is None:
|
||||
return None
|
||||
requested_logging: Final = _logging_entries(requested_metadata)
|
||||
requested_names: Final = logging_credential_names(requested_logging)
|
||||
stored_names: Final = logging_credential_names(stored_logging)
|
||||
references_involved: Final = bool(requested_names or stored_names)
|
||||
if references_involved and not is_proxy_admin and requested_logging != stored_logging:
|
||||
return CredentialReferenceRejection(
|
||||
403,
|
||||
"Only proxy admins can change the logging entries of a key that uses litellm_credential_name",
|
||||
)
|
||||
undefined: Final = sorted({name for name in set(requested_names) - set(stored_names) if not is_defined(name)})
|
||||
if undefined:
|
||||
return CredentialReferenceRejection(
|
||||
400, f"metadata.logging references undefined credential(s): {', '.join(undefined)}"
|
||||
)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -26,8 +26,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_metadata_variable_name_from_kwargs,
|
||||
get_or_create_metadata_bucket,
|
||||
)
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
|
||||
from litellm.proxy._types import CALLBACK_VAR_NAMES, CommonProxyErrors, LiteLLMPromptInjectionParams
|
||||
from litellm.proxy.common_utils.callback_config_validation import object_sequence, string_keyed
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
|
|
@ -736,6 +738,47 @@ def decrypt_callback_vars(metadata: object) -> Any:
|
|||
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
|
||||
|
||||
|
||||
def decrypted_callback_metadata(metadata: object) -> object:
|
||||
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
|
||||
|
||||
|
||||
def decrypted_logging_entries(metadata: object) -> object:
|
||||
return string_keyed(decrypted_callback_metadata(metadata)).get("logging")
|
||||
|
||||
|
||||
def key_logging_entries(metadata: object) -> tuple[Mapping[str, object], ...] | None:
|
||||
entries: Final = decrypted_logging_entries(metadata)
|
||||
if entries is None:
|
||||
return None
|
||||
return tuple(resolved for entry in _logging_objects(entries) if (resolved := _with_credential(entry)) is not None)
|
||||
|
||||
|
||||
def _logging_objects(entries: object) -> tuple[Mapping[str, object], ...]:
|
||||
items: Final = () if entries in ("", {}) else object_sequence(entries)
|
||||
if items is None or not all(isinstance(item, Mapping) for item in items):
|
||||
raise TypeError(f"key metadata.logging must be a list of objects, got {type(entries).__name__}")
|
||||
return tuple(string_keyed(item) for item in items)
|
||||
|
||||
|
||||
def _with_credential(entry: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
credential_name: Final = entry.get("litellm_credential_name")
|
||||
if not isinstance(credential_name, str):
|
||||
return entry
|
||||
credential_vars: Final = {
|
||||
name: str(value)
|
||||
for name, value in CredentialAccessor.get_credential_values(credential_name).items()
|
||||
if name in CALLBACK_VAR_NAMES and value is not None
|
||||
}
|
||||
if not credential_vars:
|
||||
verbose_proxy_logger.error(
|
||||
"Ignoring key logging entry for %s: credential %r is not defined or holds no logging values",
|
||||
entry.get("callback_name"),
|
||||
credential_name,
|
||||
)
|
||||
return None
|
||||
return {**entry, "callback_vars": {**credential_vars, **string_keyed(entry.get("callback_vars"))}}
|
||||
|
||||
|
||||
def _transform_callback_vars(metadata: object, transform: Callable[[str, object], object]) -> object:
|
||||
if not isinstance(metadata, dict):
|
||||
return metadata
|
||||
|
|
|
|||
|
|
@ -88,6 +88,7 @@ from litellm.proxy.auth.route_checks import RouteChecks
|
|||
from litellm.proxy.common_utils.callback_utils import (
|
||||
decrypt_callback_vars,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
key_logging_entries,
|
||||
strip_callback_config,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
|
||||
|
|
@ -1246,9 +1247,11 @@ def convert_key_logging_metadata_to_callback(
|
|||
return team_callback_settings_obj
|
||||
|
||||
|
||||
def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallback | None:
|
||||
def get_validated_callback_metadata(item: object, *, source: str) -> AddTeamCallback | None:
|
||||
if not isinstance(item, Mapping):
|
||||
raise TypeError(f"{source} metadata.logging must be a list of objects, got an item of {type(item).__name__}")
|
||||
try:
|
||||
return AddTeamCallback(**item)
|
||||
return AddTeamCallback.model_validate(item)
|
||||
except (PydanticValidationError, ValueError) as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid %s callback metadata: %s",
|
||||
|
|
@ -1267,9 +1270,9 @@ class KeyAndTeamLoggingSettings:
|
|||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
|
||||
def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth) -> tuple[Mapping[str, object], ...] | None:
|
||||
if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata:
|
||||
return decrypt_callback_vars(user_api_key_dict.metadata).get("logging")
|
||||
return key_logging_entries(user_api_key_dict.metadata)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1283,9 +1286,7 @@ def get_dynamic_logging_metadata(
|
|||
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
|
||||
) -> TeamCallbackMetadata | None:
|
||||
callback_settings_obj: TeamCallbackMetadata | None = None
|
||||
key_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(
|
||||
user_api_key_dict
|
||||
)
|
||||
key_dynamic_logging_settings: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
|
||||
team_dynamic_logging_settings: Final[dict | None] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(
|
||||
user_api_key_dict
|
||||
)
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq
|
|||
check_permissions_caller_permission,
|
||||
generate_key_helper_fn,
|
||||
prepare_metadata_fields,
|
||||
raise_on_logging_credential_change,
|
||||
)
|
||||
from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams
|
||||
from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports
|
||||
|
|
@ -598,6 +599,7 @@ async def new_user(
|
|||
if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth):
|
||||
enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key")
|
||||
enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key")
|
||||
raise_on_logging_credential_change(data.metadata, None, user_api_key_dict)
|
||||
|
||||
check_permissions_caller_permission(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, c
|
|||
import fastapi
|
||||
import yaml
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -68,9 +68,14 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
evict_and_broadcast,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
|
||||
from litellm.proxy.common_utils.callback_config_validation import (
|
||||
credential_reference_rejection,
|
||||
logging_metadata_config_error,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
decrypt_callback_vars,
|
||||
decrypted_callback_metadata,
|
||||
decrypted_logging_entries,
|
||||
encrypt_callback_vars,
|
||||
)
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||||
|
|
@ -266,6 +271,12 @@ class _StoredKeyRouterSettings(LiteLLMBaseModel):
|
|||
router_settings: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class _StoredKeyMetadata(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metadata: object = None
|
||||
|
||||
|
||||
class _KeyRowWhere(TypedDict):
|
||||
token: ReadOnly[str]
|
||||
|
||||
|
|
@ -934,6 +945,23 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -
|
|||
raise HTTPException(status_code=400, detail={"error": error})
|
||||
|
||||
|
||||
def raise_on_logging_credential_change(
|
||||
requested_metadata: object, existing_key_row: LiteLLM_VerificationToken | None, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> None:
|
||||
stored_metadata: Final = (
|
||||
None
|
||||
if existing_key_row is None
|
||||
else _StoredKeyMetadata.model_validate(existing_key_row, from_attributes=True).metadata
|
||||
)
|
||||
rejection: Final = credential_reference_rejection(
|
||||
decrypted_callback_metadata(requested_metadata),
|
||||
decrypted_logging_entries(stored_metadata),
|
||||
is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
if rejection is not None:
|
||||
raise HTTPException(status_code=rejection.status_code, detail={"error": rejection.message})
|
||||
|
||||
|
||||
def common_key_access_checks(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
|
|
@ -1485,6 +1513,7 @@ async def _common_key_generation_helper(
|
|||
|
||||
validate_budget_duration(data.budget_duration)
|
||||
raise_on_invalid_key_logging_config(data.metadata)
|
||||
raise_on_logging_credential_change(data.metadata, None, user_api_key_dict)
|
||||
|
||||
if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
raise HTTPException(
|
||||
|
|
@ -3430,6 +3459,7 @@ async def _process_single_key_update(
|
|||
)
|
||||
|
||||
# Prepare update data
|
||||
raise_on_logging_credential_change(key_request.metadata, existing_key_row, user_api_key_dict)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
|
@ -4176,6 +4206,7 @@ async def update_key_fn(
|
|||
|
||||
# Enforce upperbound key params on update (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
raise_on_logging_credential_change(data.metadata, existing_key_row, user_api_key_dict)
|
||||
non_default_values: Final = await prepare_key_update_data(
|
||||
data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
|
@ -6363,6 +6394,7 @@ async def _execute_virtual_key_regeneration(
|
|||
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request)
|
||||
# Enforce upperbound key params on regenerate (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
raise_on_logging_credential_change(data.metadata, key_in_db, user_api_key_dict)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
enforce_vector_store_index_grants_are_admin_only,
|
||||
)
|
||||
from litellm.proxy.auth.litellm_license import LicenseCheck
|
||||
from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
|
||||
|
|
@ -225,7 +226,14 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str
|
|||
enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key")
|
||||
except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only
|
||||
return _error_message(exc)
|
||||
return None
|
||||
rejection: Final = (
|
||||
credential_reference_rejection(
|
||||
item.metadata, None, is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
if item.auto_create_key
|
||||
else None
|
||||
)
|
||||
return None if rejection is None else rejection.message
|
||||
|
||||
|
||||
def _normalized_email(email: str | None) -> str | None:
|
||||
|
|
|
|||
|
|
@ -672,10 +672,12 @@ def _write_config(
|
|||
callback_settings: Mapping[str, JsonValue] | None = None,
|
||||
litellm_settings: Mapping[str, JsonValue] | None = None,
|
||||
general_settings: Mapping[str, JsonValue] | None = None,
|
||||
extra_config: Mapping[str, JsonValue] | None = None,
|
||||
) -> Path:
|
||||
loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config: Final = {
|
||||
**loaded,
|
||||
**(extra_config or {}),
|
||||
"litellm_settings": {
|
||||
**loaded["litellm_settings"],
|
||||
"callbacks": list(callbacks),
|
||||
|
|
@ -752,6 +754,7 @@ def _rig(
|
|||
callback_settings: Mapping[str, JsonValue] | None = None,
|
||||
litellm_settings: Mapping[str, JsonValue] | None = None,
|
||||
general_settings: Mapping[str, JsonValue] | None = None,
|
||||
extra_config: Mapping[str, JsonValue] | None = None,
|
||||
environment: Mapping[str, str] | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
disabled_environment: tuple[str, ...] = (),
|
||||
|
|
@ -782,6 +785,7 @@ def _rig(
|
|||
callback_settings=resolved_settings,
|
||||
litellm_settings=litellm_settings,
|
||||
general_settings=general_settings,
|
||||
extra_config=extra_config,
|
||||
)
|
||||
preset_environment: Final = {
|
||||
**({"PHOENIX_COLLECTOR_ENDPOINT": destination.url + "/v1/traces"} if "arize_phoenix" in callbacks else {}),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,428 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from itertools import accumulate
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from _openinference_support import (
|
||||
_anthropic_response,
|
||||
_chat_request_marker,
|
||||
_chat_response,
|
||||
_chat_stream_response,
|
||||
_chat_tool_call,
|
||||
_matching_marker_spans,
|
||||
_responses_response,
|
||||
_rig,
|
||||
)
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
CREDENTIAL_NAME: Final = "arize-space-a"
|
||||
SECRET_API_KEY: Final = "space-a-arize-key-" + uuid.uuid4().hex
|
||||
SECRET_SPACE_ID: Final = "space-a-id-" + uuid.uuid4().hex
|
||||
AMBIENT_ARIZE_ENVIRONMENT: Final = ("ARIZE_API_KEY", "ARIZE_SPACE_ID")
|
||||
CREDENTIAL_CONFIG: Final[dict[str, JsonValue]] = {
|
||||
"credential_list": [
|
||||
{
|
||||
"credential_name": CREDENTIAL_NAME,
|
||||
"credential_values": {
|
||||
"arize_api_key": "os.environ/SPACE_A_ARIZE_API_KEY",
|
||||
"arize_space_id": "os.environ/SPACE_A_ARIZE_SPACE_ID",
|
||||
},
|
||||
"credential_info": {},
|
||||
},
|
||||
{
|
||||
"credential_name": "model-only",
|
||||
"credential_values": {"api_key": "sk-not-a-logging-value"},
|
||||
"credential_info": {},
|
||||
},
|
||||
]
|
||||
}
|
||||
SECRET_ENVIRONMENT: Final = {"SPACE_A_ARIZE_API_KEY": SECRET_API_KEY, "SPACE_A_ARIZE_SPACE_ID": SECRET_SPACE_ID}
|
||||
|
||||
|
||||
def _reference_entry(credential_name: str = CREDENTIAL_NAME) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"litellm_credential_name": credential_name,
|
||||
"callback_vars": {},
|
||||
}
|
||||
|
||||
|
||||
def _echo_marker(request: Request) -> Reply:
|
||||
return _chat_response(_chat_request_marker(request))
|
||||
|
||||
|
||||
def _send_chat(proxy: Gateway, model: str, key: str, marker: str) -> httpx.Response:
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"metadata": {"trace_marker": marker},
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _chat(proxy: Gateway, model: str, key: str, marker: str) -> None:
|
||||
response: Final = _send_chat(proxy, model, key, marker)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _drains(destination: Wire) -> Iterator[tuple[Request, ...]]:
|
||||
while True:
|
||||
yield destination.drain()
|
||||
time.sleep(0.1)
|
||||
|
||||
|
||||
def _carries(request: Request, marker: str) -> bool:
|
||||
return any(True for _ in _matching_marker_spans((request,), marker))
|
||||
|
||||
|
||||
def _received_until(destination: Wire, marker: str, seconds: float = 30) -> tuple[Request, ...]:
|
||||
deadline: Final = time.monotonic() + seconds
|
||||
received: Final = next(
|
||||
batch
|
||||
for batch in accumulate(_drains(destination), operator.add)
|
||||
if any(_carries(request, marker) for request in batch) or time.monotonic() >= deadline
|
||||
)
|
||||
assert any(_carries(request, marker) for request in received), f"no span for {marker} reached the destination"
|
||||
return received
|
||||
|
||||
|
||||
def test_key_referencing_a_credential_exports_with_the_secret_manager_values(gateway: Gateway, tmp_path: Path) -> None:
|
||||
unreferenced_marker: Final = "plain-" + uuid.uuid4().hex
|
||||
referenced_marker: Final = "ref-" + uuid.uuid4().hex
|
||||
with (
|
||||
_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
_echo_marker,
|
||||
extra_config=CREDENTIAL_CONFIG,
|
||||
environment=SECRET_ENVIRONMENT,
|
||||
remove_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
disabled_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
team: Final = scenario.team()
|
||||
referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry()]})
|
||||
unreferenced_key: Final = scenario.key(team_id=team)
|
||||
|
||||
_chat(rig.proxy, rig.model, unreferenced_key, unreferenced_marker)
|
||||
_chat(rig.proxy, rig.model, referencing_key, referenced_marker)
|
||||
received: Final = _received_until(rig.destination, referenced_marker)
|
||||
|
||||
exports: Final = tuple(request for request in received if _carries(request, referenced_marker))
|
||||
for export in exports:
|
||||
headers = {name.lower(): value for name, value in export.headers.items()}
|
||||
assert headers.get("api_key") == SECRET_API_KEY, headers
|
||||
assert headers.get("arize-space-id") == SECRET_SPACE_ID, headers
|
||||
assert not any(_carries(request, unreferenced_marker) for request in received)
|
||||
|
||||
info: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key})
|
||||
assert info.status_code == 200, info.text
|
||||
assert CREDENTIAL_NAME in info.text
|
||||
assert SECRET_API_KEY not in info.text
|
||||
assert SECRET_SPACE_ID not in info.text
|
||||
|
||||
|
||||
def test_only_proxy_admins_change_logging_on_a_key_that_references_a_credential(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
_rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
team: Final = scenario.team()
|
||||
team_admin: Final = scenario.member(team, role="admin")
|
||||
team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team)
|
||||
langfuse_entry: Final[dict[str, JsonValue]] = {
|
||||
"callback_name": "langfuse",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"langfuse_public_key": "pk-lf-integration", "langfuse_secret_key": "sk-lf-integration"},
|
||||
}
|
||||
referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry(), langfuse_entry]})
|
||||
|
||||
def as_team_admin(path: str, body: dict[str, JsonValue]) -> httpx.Response:
|
||||
return rig.proxy.request("POST", path, body, key=team_admin_key)
|
||||
|
||||
def assert_gated(response: httpx.Response) -> None:
|
||||
assert response.status_code == 403, response.text
|
||||
assert "litellm_credential_name" in response.text, response.text
|
||||
|
||||
assert_gated(as_team_admin("/key/generate", {"team_id": team, "metadata": {"logging": [_reference_entry()]}}))
|
||||
|
||||
sibling: Final = as_team_admin(
|
||||
"/key/update",
|
||||
{
|
||||
"key": referencing_key,
|
||||
"metadata": {
|
||||
"logging": [
|
||||
_reference_entry(),
|
||||
{
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"arize_endpoint": "http://127.0.0.1:9/v1/traces"},
|
||||
},
|
||||
]
|
||||
},
|
||||
},
|
||||
)
|
||||
assert_gated(sibling)
|
||||
assert_gated(as_team_admin("/key/update", {"key": referencing_key, "metadata": {"team": "unrelated"}}))
|
||||
assert_gated(as_team_admin("/key/regenerate", {"key": referencing_key, "metadata": {"team": "unrelated"}}))
|
||||
|
||||
bulk: Final = as_team_admin(
|
||||
"/team/key/bulk_update",
|
||||
{"team_id": team, "key_ids": [referencing_key], "update_fields": {"metadata": {"team": "unrelated"}}},
|
||||
)
|
||||
assert bulk.status_code == 200, bulk.text
|
||||
assert bulk.json()["successful_updates"] == [], bulk.text
|
||||
assert "litellm_credential_name" in bulk.json()["failed_updates"][0]["failed_reason"], bulk.text
|
||||
|
||||
renamed: Final = as_team_admin(
|
||||
"/key/update", {"key": referencing_key, "key_alias": "alias-" + uuid.uuid4().hex}
|
||||
)
|
||||
assert renamed.status_code == 200, renamed.text
|
||||
|
||||
read_back: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}, key=team_admin_key)
|
||||
assert read_back.status_code == 200, read_back.text
|
||||
resaved: Final = as_team_admin(
|
||||
"/key/update",
|
||||
{"key": referencing_key, "metadata": {**read_back.json()["info"]["metadata"], "note": "unrelated edit"}},
|
||||
)
|
||||
assert resaved.status_code == 200, resaved.text
|
||||
|
||||
stored: Final = rig.proxy.get("/key/info", params={"key": referencing_key})
|
||||
logging_entries: Final = stored["info"]["metadata"]["logging"]
|
||||
assert [entry["callback_name"] for entry in logging_entries] == ["arize", "langfuse"], stored
|
||||
assert logging_entries[0] == _reference_entry(), stored
|
||||
|
||||
undefined: Final = rig.proxy.request(
|
||||
"POST", "/key/generate", {"metadata": {"logging": [_reference_entry("no-such-credential")]}}
|
||||
)
|
||||
assert undefined.status_code == 400, undefined.text
|
||||
assert "no-such-credential" in undefined.text, undefined.text
|
||||
|
||||
|
||||
def test_only_proxy_admins_create_a_user_whose_key_references_a_credential(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
_rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
organization: Final = scenario.organization()
|
||||
org_admin_key: Final = scenario.key(user_id=scenario.org_member(organization, "org_admin"))
|
||||
referencing_metadata: Final = {"logging": [_reference_entry()]}
|
||||
|
||||
single: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/user/new",
|
||||
{"user_id": "user-" + uuid.uuid4().hex, "organizations": [organization], "metadata": referencing_metadata},
|
||||
key=org_admin_key,
|
||||
)
|
||||
assert single.status_code == 403, single.text
|
||||
assert "litellm_credential_name" in single.text, single.text
|
||||
|
||||
delegate_key: Final = scenario.key(
|
||||
user_id=scenario.user(user_role="internal_user"), allowed_routes=["/management/v1/users/bulk"]
|
||||
)
|
||||
bulk: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/management/v1/users/bulk",
|
||||
{
|
||||
"users": [
|
||||
{"user_id": "user-" + uuid.uuid4().hex, "auto_create_key": True, "metadata": referencing_metadata}
|
||||
]
|
||||
},
|
||||
key=delegate_key,
|
||||
)
|
||||
assert bulk.status_code == 200, bulk.text
|
||||
assert [row["success"] for row in bulk.json()["data"]] == [False], bulk.text
|
||||
assert "litellm_credential_name" in bulk.json()["data"][0]["error"], bulk.text
|
||||
|
||||
|
||||
def _carries_credential(request: Request) -> bool:
|
||||
headers: Final = {name.lower(): value for name, value in request.headers.items()}
|
||||
return headers.get("api_key") == SECRET_API_KEY and headers.get("arize-space-id") == SECRET_SPACE_ID
|
||||
|
||||
|
||||
def _credential_exports(destination: Wire, seconds: float = 30) -> tuple[Request, ...]:
|
||||
return eventually(
|
||||
lambda: tuple(request for request in destination.drain() if _carries_credential(request)),
|
||||
bool,
|
||||
seconds=seconds,
|
||||
)
|
||||
|
||||
|
||||
def _surface_reply(marker: str) -> Callable[[Request], Reply]:
|
||||
def reply(request: Request) -> Reply:
|
||||
if request.target.endswith("/responses"):
|
||||
return _responses_response(marker)
|
||||
if request.target.endswith("/messages"):
|
||||
return _anthropic_response(marker)
|
||||
return _chat_stream_response(marker, (_chat_tool_call(marker),))
|
||||
|
||||
return reply
|
||||
|
||||
|
||||
_SURFACES: Final[dict[str, tuple[str, str, dict[str, JsonValue]]]] = {
|
||||
"responses": ("/v1/responses", "openai/gpt-4o-mini", {"input": "weather in Paris?"}),
|
||||
"messages": (
|
||||
"/v1/messages",
|
||||
"anthropic/claude-opus-5-5",
|
||||
{"max_tokens": 64, "messages": [{"role": "user", "content": "weather in Paris?"}]},
|
||||
),
|
||||
"chat-stream": (
|
||||
"/v1/chat/completions",
|
||||
"openai/gpt-4o-mini",
|
||||
{"stream": True, "messages": [{"role": "user", "content": "weather in Paris?"}]},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", sorted(_SURFACES))
|
||||
def test_referencing_key_exports_on_each_surface(surface: str, gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = f"{surface}-" + uuid.uuid4().hex
|
||||
path, model_name, body = _SURFACES[surface]
|
||||
with (
|
||||
_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
_surface_reply(marker),
|
||||
extra_config=CREDENTIAL_CONFIG,
|
||||
environment=SECRET_ENVIRONMENT,
|
||||
remove_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
disabled_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
model_name=model_name,
|
||||
api_base_suffix="" if surface == "messages" else "/v1",
|
||||
) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key(metadata={"logging": [_reference_entry()]})
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
path,
|
||||
{"model": rig.model, **body, "metadata": {"trace_marker": marker}, "cache": {"no-cache": True}},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
received: Final = _received_until(rig.destination, marker)
|
||||
exports: Final = tuple(request for request in received if _carries(request, marker))
|
||||
assert exports and all(_carries_credential(request) for request in exports), [
|
||||
dict(request.headers) for request in exports
|
||||
]
|
||||
|
||||
|
||||
def test_referencing_key_exports_through_anthropic_pass_through(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "passthrough-" + uuid.uuid4().hex
|
||||
with wire_server(lambda _request: _anthropic_response(marker)) as anthropic:
|
||||
with (
|
||||
_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
_echo_marker,
|
||||
extra_config=CREDENTIAL_CONFIG,
|
||||
environment={
|
||||
**SECRET_ENVIRONMENT,
|
||||
"ANTHROPIC_API_BASE": anthropic.url,
|
||||
"ANTHROPIC_API_KEY": "sk-ant-synthetic",
|
||||
},
|
||||
remove_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
disabled_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key(metadata={"logging": [_reference_entry()]})
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/anthropic/v1/messages",
|
||||
{"model": "claude-opus-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["id"] == marker, response.text
|
||||
assert anthropic.drain(), "the pass-through never reached the scripted Anthropic upstream"
|
||||
assert _credential_exports(rig.destination)
|
||||
|
||||
|
||||
def test_referencing_key_exports_on_otel_v1(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "otel-v1-" + uuid.uuid4().hex
|
||||
with (
|
||||
_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
_echo_marker,
|
||||
extra_config=CREDENTIAL_CONFIG,
|
||||
environment={**SECRET_ENVIRONMENT, "LITELLM_OTEL_V2": "false"},
|
||||
remove_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
disabled_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key(metadata={"logging": [_reference_entry()]})
|
||||
_chat(rig.proxy, rig.model, key, marker)
|
||||
assert _credential_exports(rig.destination)
|
||||
|
||||
|
||||
def test_reference_to_a_credential_without_logging_values_exports_nothing_and_is_unhealthy(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
dropped_marker: Final = "model-only-" + uuid.uuid4().hex
|
||||
sentinel_marker: Final = "sentinel-" + uuid.uuid4().hex
|
||||
with (
|
||||
_rig(
|
||||
gateway,
|
||||
tmp_path,
|
||||
_echo_marker,
|
||||
extra_config=CREDENTIAL_CONFIG,
|
||||
environment=SECRET_ENVIRONMENT,
|
||||
remove_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
disabled_environment=AMBIENT_ARIZE_ENVIRONMENT,
|
||||
) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
dropped_key: Final = scenario.key(metadata={"logging": [_reference_entry("model-only")]})
|
||||
sentinel_key: Final = scenario.key(metadata={"logging": [_reference_entry()]})
|
||||
|
||||
_chat(rig.proxy, rig.model, dropped_key, dropped_marker)
|
||||
_chat(rig.proxy, rig.model, sentinel_key, sentinel_marker)
|
||||
received: Final = _received_until(rig.destination, sentinel_marker)
|
||||
assert not any(_carries(request, dropped_marker) for request in received)
|
||||
|
||||
health: Final = rig.proxy.request("POST", "/key/health", {}, key=dropped_key)
|
||||
assert health.status_code == 200, health.text
|
||||
assert health.json()["key"] == "unhealthy", health.text
|
||||
|
||||
|
||||
def test_malformed_key_logging_fails_requests_like_main(gateway: Gateway, tmp_path: Path) -> None:
|
||||
stored_logging: Final[dict[str, JsonValue]] = {
|
||||
"string": "abc",
|
||||
"list-with-string": ["x"],
|
||||
"object": {"callback_name": "arize"},
|
||||
"empty-list": [],
|
||||
}
|
||||
with (
|
||||
_rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig,
|
||||
rig.proxy.scenario() as scenario,
|
||||
):
|
||||
statuses: Final = {
|
||||
label: _send_chat(
|
||||
rig.proxy, rig.model, scenario.key(metadata={"logging": value}), "malformed-" + uuid.uuid4().hex
|
||||
).status_code
|
||||
for label, value in stored_logging.items()
|
||||
}
|
||||
|
||||
assert statuses == {"string": 500, "list-with-string": 500, "object": 500, "empty-list": 200}, statuses
|
||||
|
|
@ -7,6 +7,7 @@ from litellm.proxy.common_utils.callback_config_validation import (
|
|||
callback_config_error,
|
||||
conflicting_callback_var_error,
|
||||
conflicting_span_scope_error,
|
||||
credential_reference_rejection,
|
||||
cross_entry_family_error,
|
||||
logging_metadata_config_error,
|
||||
)
|
||||
|
|
@ -499,3 +500,78 @@ def test_key_logging_entries_compare_excluded_services_by_parsed_value(
|
|||
def test_a_failure_only_entry_rejects_excluded_services() -> None:
|
||||
error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}, "failure")
|
||||
assert error is not None and "excluded_services" in error and "success_and_failure" in error
|
||||
|
||||
|
||||
_REFERENCE: Final = {"callback_name": "arize", "callback_type": "success", "litellm_credential_name": "space-a"}
|
||||
_INLINE: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}
|
||||
_SIBLING: Final = {"callback_name": "arize", "callback_vars": {"arize_endpoint": "http://attacker.example/v1/traces"}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("requested", "stored_logging", "is_proxy_admin", "expected_status"),
|
||||
[
|
||||
pytest.param({"logging": [_REFERENCE]}, None, False, 403, id="non-admin-adds-reference"),
|
||||
pytest.param({"logging": [_INLINE, _REFERENCE]}, None, False, 403, id="non-admin-adds-reference-second"),
|
||||
pytest.param({"team": "x"}, [_REFERENCE], False, 403, id="non-admin-removes-reference"),
|
||||
pytest.param({"logging": [_REFERENCE, _SIBLING]}, [_REFERENCE], False, 403, id="non-admin-adds-sibling"),
|
||||
pytest.param({"logging": [_REFERENCE]}, [_REFERENCE], False, None, id="non-admin-resends-same"),
|
||||
pytest.param(None, [_REFERENCE], False, None, id="non-admin-leaves-metadata-alone"),
|
||||
pytest.param({"logging": [_INLINE]}, [], False, None, id="non-admin-without-references"),
|
||||
pytest.param({"logging": [_REFERENCE]}, [], True, None, id="admin-adds-defined-reference"),
|
||||
pytest.param({"team": "x"}, [_REFERENCE], True, None, id="admin-removes-reference"),
|
||||
],
|
||||
)
|
||||
def test_credential_reference_rejection_lets_only_proxy_admins_change_referencing_logging(
|
||||
requested: dict[str, object] | None,
|
||||
stored_logging: list[object] | None,
|
||||
is_proxy_admin: bool,
|
||||
expected_status: int | None,
|
||||
) -> None:
|
||||
rejection: Final = credential_reference_rejection(
|
||||
requested, stored_logging, is_proxy_admin=is_proxy_admin, is_defined=lambda name: name == "space-a"
|
||||
)
|
||||
assert (None if rejection is None else rejection.status_code) == expected_status
|
||||
|
||||
|
||||
def test_credential_reference_rejection_names_undefined_credentials_only_to_admins() -> None:
|
||||
requested: Final = {"logging": [_INLINE, {**_REFERENCE, "litellm_credential_name": "missing"}]}
|
||||
|
||||
admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=True, is_defined=lambda _: False)
|
||||
non_admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=False, is_defined=lambda _: False)
|
||||
|
||||
assert admin is not None and admin.status_code == 400 and "missing" in admin.message
|
||||
assert non_admin is not None and non_admin.status_code == 403 and "missing" not in non_admin.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_proxy_admin", [True, False], ids=["admin", "non-admin"])
|
||||
def test_credential_reference_rejection_lets_a_stored_reference_to_a_deleted_credential_stay(
|
||||
is_proxy_admin: bool,
|
||||
) -> None:
|
||||
stale: Final = {**_REFERENCE, "litellm_credential_name": "deleted"}
|
||||
|
||||
kept: Final = credential_reference_rejection(
|
||||
{"logging": [stale], "note": "unrelated edit"},
|
||||
[stale],
|
||||
is_proxy_admin=is_proxy_admin,
|
||||
is_defined=lambda _: False,
|
||||
)
|
||||
added: Final = credential_reference_rejection(
|
||||
{"logging": [stale, {**_REFERENCE, "litellm_credential_name": "also-missing"}]},
|
||||
[stale],
|
||||
is_proxy_admin=True,
|
||||
is_defined=lambda _: False,
|
||||
)
|
||||
|
||||
assert kept is None
|
||||
assert added is not None and added.status_code == 400
|
||||
assert "also-missing" in added.message and "deleted" not in added.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", [5, ["space-a"], None], ids=["int", "list", "null"])
|
||||
def test_a_non_string_credential_name_is_not_a_reference(name: object) -> None:
|
||||
entry: Final = {**_REFERENCE, "litellm_credential_name": name}
|
||||
|
||||
assert (
|
||||
credential_reference_rejection({"logging": [entry]}, None, is_proxy_admin=False, is_defined=lambda _: False)
|
||||
is None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import copy
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Final
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
initialize_callbacks_on_proxy,
|
||||
key_logging_entries,
|
||||
normalize_callback_names,
|
||||
process_callback,
|
||||
sanitize_openai_provider_metadata,
|
||||
|
|
@ -28,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
def test_get_remaining_tokens_and_requests_from_request_data():
|
||||
|
|
@ -775,3 +778,114 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con
|
|||
|
||||
assert len(litellm.callbacks) == 1
|
||||
assert isinstance(litellm.callbacks[0], CustomLogger)
|
||||
|
||||
|
||||
def _credential(name: str, values: dict[str, object]) -> CredentialItem:
|
||||
return CredentialItem(credential_name=name, credential_values=values, credential_info={})
|
||||
|
||||
|
||||
def test_key_logging_entries_merges_each_referenced_credential_under_its_entry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
_credential("space-a", {"arize_api_key": "key-a", "arize_space_id": "space-a", "api_base": "x"}),
|
||||
_credential("space-b", {"arize_api_key": "key-b", "arize_space_id": "space-b"}),
|
||||
],
|
||||
)
|
||||
entries: Final = [
|
||||
{"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}},
|
||||
{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {"arize_space_id": "own"}},
|
||||
{"callback_name": "arize", "litellm_credential_name": "space-b", "callback_vars": {}},
|
||||
]
|
||||
|
||||
resolved: Final = key_logging_entries({"logging": entries})
|
||||
|
||||
assert resolved is not None
|
||||
assert [entry["callback_vars"] for entry in resolved] == [
|
||||
{"langfuse_public_key": "pk"},
|
||||
{"arize_api_key": "key-a", "arize_space_id": "own"},
|
||||
{"arize_api_key": "key-b", "arize_space_id": "space-b"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"credential_values",
|
||||
[{"api_key": "sk-model-key"}, {"arize_api_key": None}],
|
||||
ids=["model-credential", "null-values"],
|
||||
)
|
||||
def test_key_logging_entries_drops_an_entry_whose_credential_supplies_no_logging_values(
|
||||
monkeypatch: pytest.MonkeyPatch, credential_values: dict[str, object]
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential("model-only", credential_values)])
|
||||
kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}
|
||||
entries: Final = [
|
||||
kept,
|
||||
{"callback_name": "arize", "litellm_credential_name": "model-only", "callback_vars": {}},
|
||||
]
|
||||
|
||||
assert key_logging_entries({"logging": entries}) == (kept,)
|
||||
|
||||
|
||||
def test_a_dropped_credential_reference_is_logged_as_an_error_so_key_health_reports_it(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [])
|
||||
entry: Final = {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}}
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
key_logging_entries({"logging": [entry]})
|
||||
|
||||
assert [record.levelno for record in caplog.records if "missing" in record.getMessage()] == [logging.ERROR]
|
||||
|
||||
|
||||
def test_key_logging_entries_drops_an_entry_whose_credential_is_not_defined(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential("space-a", {"arize_api_key": "key-a"})])
|
||||
kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}
|
||||
entries: Final = [
|
||||
kept,
|
||||
{"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}},
|
||||
]
|
||||
|
||||
assert key_logging_entries({"logging": entries}) == (kept,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"logging_value",
|
||||
[
|
||||
pytest.param("abc", id="string"),
|
||||
pytest.param({"callback_name": "arize"}, id="object"),
|
||||
pytest.param(5, id="number"),
|
||||
pytest.param(["x"], id="list-with-string"),
|
||||
pytest.param([{"callback_name": "langfuse", "callback_vars": {}}, None], id="list-with-null"),
|
||||
],
|
||||
)
|
||||
def test_key_logging_entries_fails_on_non_object_logging_like_the_request_path_always_has(
|
||||
logging_value: object,
|
||||
) -> None:
|
||||
with pytest.raises(TypeError):
|
||||
key_logging_entries({"logging": logging_value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("logging_value", "expected"),
|
||||
[
|
||||
pytest.param("", (), id="empty-string"),
|
||||
pytest.param({}, (), id="empty-object"),
|
||||
pytest.param([], (), id="empty-list"),
|
||||
pytest.param([{"bad": 1}], ({"bad": 1},), id="object-with-unknown-fields"),
|
||||
],
|
||||
)
|
||||
def test_key_logging_entries_passes_empty_and_object_entries_through(
|
||||
logging_value: object, expected: tuple[dict[str, object], ...]
|
||||
) -> None:
|
||||
assert key_logging_entries({"logging": logging_value}) == expected
|
||||
|
||||
|
||||
def test_key_logging_entries_leaves_an_entry_with_a_non_string_credential_name_as_written(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential("5", {"arize_api_key": "k"})])
|
||||
entry: Final = {"callback_name": "arize", "litellm_credential_name": 5, "callback_vars": {}}
|
||||
|
||||
assert key_logging_entries({"logging": [entry]}) == (entry,)
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
|
||||
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
budget_windows_absent_cache_key,
|
||||
|
|
@ -101,6 +102,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
persist_deleted_verification_tokens,
|
||||
prepare_key_update_data,
|
||||
prepare_metadata_fields,
|
||||
raise_on_logging_credential_change,
|
||||
reset_key_spend_fn,
|
||||
update_key_fn,
|
||||
validate_key_list_check,
|
||||
|
|
@ -112,6 +114,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
|||
BulkUpdateKeyResponse,
|
||||
CustomKeyPolicyRequest,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
client = TestClient(app)
|
||||
|
|
@ -24407,3 +24410,29 @@ async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypa
|
|||
)
|
||||
|
||||
assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection")
|
||||
|
||||
|
||||
def test_non_admin_resaving_referencing_key_metadata_read_back_encrypted_is_not_a_logging_change(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-32-bytes-aaaaaaaaaaaaaa")
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})],
|
||||
)
|
||||
stored = encrypt_callback_vars(
|
||||
{
|
||||
"logging": [
|
||||
{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}},
|
||||
{"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk-lf-secret"}},
|
||||
]
|
||||
}
|
||||
)
|
||||
assert stored["logging"][1]["callback_vars"]["langfuse_secret_key"] != "sk-lf-secret"
|
||||
key_row = LiteLLM_VerificationToken(token="hashed", metadata=stored)
|
||||
internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
|
||||
raise_on_logging_credential_change({**stored, "note": "unrelated edit"}, key_row, internal_user)
|
||||
with pytest.raises(HTTPException) as removed:
|
||||
raise_on_logging_credential_change({"note": "unrelated edit"}, key_row, internal_user)
|
||||
|
||||
assert removed.value.status_code == 403
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import pytest
|
|||
from prisma.errors import UniqueViolationError
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth
|
||||
from litellm.proxy.list_api.common import ManagementProblem
|
||||
|
|
@ -15,6 +16,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import (
|
|||
BulkNewUserItem,
|
||||
BulkNewUserRequest,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
|
|
@ -434,6 +436,39 @@ async def test_only_proxy_admin_sets_batch_limits_on_a_created_key():
|
|||
assert calls[-1]["metadata"] == limits
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_proxy_admin_creates_a_key_whose_logging_references_a_credential(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})],
|
||||
)
|
||||
reference: Final = {"logging": [{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}]}
|
||||
calls: list[dict[str, object]] = []
|
||||
|
||||
async def generate_key(**kwargs: object) -> dict[str, object]:
|
||||
calls.append(kwargs)
|
||||
return {"token": f"sk-{kwargs['user_id']}"}
|
||||
|
||||
rows: Final = [
|
||||
{"user_id": "u1", "auto_create_key": True, "metadata": reference},
|
||||
{"user_id": "u2", "auto_create_key": True},
|
||||
]
|
||||
|
||||
response = await _run(_FakePrisma(), rows, caller=INTERNAL, generate_key=generate_key)
|
||||
|
||||
assert [r.success for r in response.data] == [False, True]
|
||||
assert "litellm_credential_name" in (response.data[0].error or "")
|
||||
assert [call["user_id"] for call in calls] == ["u2"]
|
||||
|
||||
admin_response = await _run(_FakePrisma(), rows[:1], caller=ADMIN, generate_key=generate_key)
|
||||
|
||||
assert [r.key for r in admin_response.data] == ["sk-u1"]
|
||||
assert calls[-1]["metadata"] == reference
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_license_is_checked_once_against_the_whole_batch():
|
||||
prisma = _FakePrisma()
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm.constants import (
|
|||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
)
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import (
|
||||
ProviderSpecificHeaderUtils,
|
||||
|
|
@ -59,6 +60,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
get_dynamic_logging_metadata,
|
||||
get_metadata_variable_name,
|
||||
move_guardrails_to_metadata,
|
||||
resolve_tenant_otel_destinations,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
|
||||
|
|
@ -2779,7 +2781,7 @@ def test_key_dynamic_logging_settings():
|
|||
team_metadata={},
|
||||
)
|
||||
result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize)
|
||||
assert result == [{"callback_name": "arize", "callback_type": "success"}]
|
||||
assert result == ({"callback_name": "arize", "callback_type": "success"},)
|
||||
|
||||
# Test with langfuse logging
|
||||
key_with_langfuse = UserAPIKeyAuth(
|
||||
|
|
@ -2788,7 +2790,7 @@ def test_key_dynamic_logging_settings():
|
|||
team_metadata={},
|
||||
)
|
||||
result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse)
|
||||
assert result == [{"callback_name": "langfuse", "callback_type": "success"}]
|
||||
assert result == ({"callback_name": "langfuse", "callback_type": "success"},)
|
||||
|
||||
# Test with no logging metadata
|
||||
key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={})
|
||||
|
|
@ -2912,6 +2914,22 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging():
|
|||
assert result.callback_vars["arize_space_id"] == "test_arize_space_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"logging_value",
|
||||
[
|
||||
pytest.param({"callback_name": "langsmith"}, id="object"),
|
||||
pytest.param("langsmith", id="string"),
|
||||
pytest.param(["langsmith"], id="list-with-string"),
|
||||
pytest.param([None], id="list-with-null"),
|
||||
],
|
||||
)
|
||||
def test_non_object_team_logging_metadata_fails_as_a_server_error_not_a_bad_request(logging_value: object) -> None:
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={"logging": logging_value})
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock())
|
||||
|
||||
|
||||
def test_add_team_callback_rejects_env_reference():
|
||||
with pytest.raises(PydanticValidationError) as exc_info:
|
||||
AddTeamCallback(
|
||||
|
|
@ -10267,6 +10285,48 @@ def test_arize_otlp_protocol_on_a_key_logging_entry_reaches_the_destination(monk
|
|||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
|
||||
def test_a_key_referencing_a_credential_routes_to_that_credential_with_its_own_excluded_services(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "https://arize.internal.example/v1/traces")
|
||||
monkeypatch.delenv("ARIZE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ARIZE_SPACE_ID", raising=False)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="space-a",
|
||||
credential_values={"arize_api_key": "k", "arize_space_id": "s"},
|
||||
credential_info={},
|
||||
)
|
||||
],
|
||||
)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
metadata={
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"litellm_credential_name": "space-a",
|
||||
"callback_vars": {"excluded_services": "redis,postgres"},
|
||||
}
|
||||
]
|
||||
},
|
||||
team_metadata={},
|
||||
)
|
||||
|
||||
destinations = resolve_tenant_otel_destinations(auth)
|
||||
|
||||
assert [(d.headers, d.excluded_db_systems) for d in destinations] == [
|
||||
({"arize-space-id": "s", "api_key": "k"}, frozenset({"redis", "postgresql"}))
|
||||
]
|
||||
finally:
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
|
||||
class TestResolveUserProviderCredentials:
|
||||
"""``_resolve_user_provider_credentials_for_request`` injects the calling user's stored
|
||||
provider tokens into secret_fields for deployments carrying a per_user_oauth credential."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue