mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(proxy): let a key reference a callback_settings definition by name
A callback_settings entry with a callback_type other than generic_api is a
named definition. A key's logging entry naming it runs that backend with the
definition's values, which the config loader resolved through the secret
manager at boot, so the key stores no credential.
A reference entry carries the name and event type only, and is the only entry
on the key for its backend, so no key-side value can merge with the
definition's credentials. Only a proxy admin may point a key at a definition
or change the logging list of a key that holds one, on /key/generate,
/key/update, /key/{key}/regenerate and the metadata of /user/new. Boot fails
on a definition named after a callback, with an unknown callback_type, a
field no callback accepts, or a secret the secret manager cannot provide,
naming the field and never its value
This commit is contained in:
Jesus Velarde
parent
bbc7de1ccb
commit
e44e42dfe4
10 changed files with 862 additions and 9 deletions
|
|
@ -6,10 +6,14 @@ per-integration checks here.
|
|||
"""
|
||||
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterator, 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, TypeAlias, cast # noqa: TID251 # preserve arbitrary keys from validated callback mappings
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR, EXCLUDED_SERVICES_VAR
|
||||
|
||||
_NEWRELIC_CALLBACK: Final = "newrelic"
|
||||
|
|
@ -21,6 +25,148 @@ _ARIZE_OTLP_PROTOCOL_VAR: Final = "arize_otlp_protocol"
|
|||
_ARIZE_SAMPLING_RATE_VARS: Final[frozenset[str]] = frozenset(
|
||||
{"arize_success_sampling_rate", "arize_error_sampling_rate"}
|
||||
)
|
||||
_GENERIC_API_CALLBACK_TYPE: Final = "generic_api"
|
||||
_CALLBACK_TYPE_FIELD: Final = "callback_type"
|
||||
_SECRET_REFERENCE_PREFIX: Final = "os.environ/"
|
||||
_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
LoggingEntryKey: TypeAlias = tuple[str | None, str | None, tuple[tuple[str, str], ...]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CallbackDefinition:
|
||||
"""A ``callback_settings`` entry a key's ``logging`` entry may reference by name.
|
||||
|
||||
``backend`` is the callback it runs on and ``callback_vars`` the values the config
|
||||
loader resolved at boot, secrets included, so the key stores only the name.
|
||||
"""
|
||||
|
||||
name: str
|
||||
backend: str
|
||||
callback_vars: Mapping[str, str]
|
||||
|
||||
|
||||
def callback_definition(callback_name: object) -> CallbackDefinition | None:
|
||||
"""The definition ``callback_name`` references, or None when it is a plain callback name."""
|
||||
if not isinstance(callback_name, str):
|
||||
return None
|
||||
return _parsed_definition(callback_name, litellm.callback_settings.get(callback_name))
|
||||
|
||||
|
||||
def _parsed_definition(name: str, raw: object) -> CallbackDefinition | None:
|
||||
if not isinstance(raw, Mapping):
|
||||
return None
|
||||
definition: Final = _STRING_KEYED_MAPPING.validate_python(raw)
|
||||
backend: Final = definition.get(_CALLBACK_TYPE_FIELD)
|
||||
if not isinstance(backend, str) or backend == _GENERIC_API_CALLBACK_TYPE:
|
||||
return None
|
||||
return CallbackDefinition(
|
||||
name=name,
|
||||
backend=backend,
|
||||
callback_vars=MappingProxyType(
|
||||
{
|
||||
key: "" if value is None else str(value)
|
||||
for key, value in definition.items()
|
||||
if key != _CALLBACK_TYPE_FIELD
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def callback_backend_name(callback_name: str) -> str:
|
||||
"""The callback a ``logging`` entry runs on: its definition's backend for a reference, else its own name."""
|
||||
definition: Final = callback_definition(callback_name)
|
||||
return callback_name if definition is None else definition.backend
|
||||
|
||||
|
||||
def callback_definitions_error(callback_settings: Mapping[str, object]) -> str | None:
|
||||
"""Why a ``callback_settings`` definition could never serve a key, checked at config load.
|
||||
|
||||
A typo in a field or secret name would otherwise surface only on the first request
|
||||
of a key that already depends on the definition. Errors name the field, never its value.
|
||||
"""
|
||||
definitions: Final = (
|
||||
definition
|
||||
for name, raw in callback_settings.items()
|
||||
if (definition := _parsed_definition(name, raw)) is not None
|
||||
)
|
||||
return next((error for definition in definitions if (error := _definition_error(definition)) is not None), None)
|
||||
|
||||
|
||||
def _definition_error(definition: CallbackDefinition) -> str | None:
|
||||
known_callbacks: Final = litellm._known_custom_logger_compatible_callbacks # pyright: ignore[reportPrivateUsage] # the registry of callback names has no public accessor
|
||||
if definition.name in known_callbacks:
|
||||
return (
|
||||
f"callback_settings.{definition.name} is named after the {definition.name} callback, "
|
||||
"so a key could not tell them apart"
|
||||
)
|
||||
if definition.backend not in known_callbacks:
|
||||
return f"callback_settings.{definition.name}.callback_type {definition.backend!r} is not a known callback"
|
||||
unresolved: Final = next(
|
||||
(
|
||||
key
|
||||
for key, value in definition.callback_vars.items()
|
||||
if value == "" or value.startswith(_SECRET_REFERENCE_PREFIX)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if unresolved is not None:
|
||||
return (
|
||||
f"callback_settings.{definition.name}.{unresolved} is empty or could not be resolved "
|
||||
"from the secret manager"
|
||||
)
|
||||
return _definition_backend_error(definition)
|
||||
|
||||
|
||||
def _definition_backend_error(definition: CallbackDefinition) -> str | None:
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.proxy._types import AddTeamCallback
|
||||
|
||||
try:
|
||||
validated: Final = AddTeamCallback.model_validate(
|
||||
{"callback_name": definition.backend, "callback_vars": dict(definition.callback_vars)}
|
||||
)
|
||||
except ValidationError as e:
|
||||
return f"callback_settings.{definition.name}: " + "; ".join(error["msg"] for error in e.errors())
|
||||
error: Final = callback_config_error(definition.backend, validated.callback_vars)
|
||||
return None if error is None else f"callback_settings.{definition.name}: {error}"
|
||||
|
||||
|
||||
def callback_reference_error(logging_entries: object) -> str | None:
|
||||
"""A ``callback_settings`` reference is a name only, and the only entry for its backend.
|
||||
|
||||
Every entry's vars merge into one map per backend at request time, so a value on the
|
||||
reference or on a sibling naming the same backend could point the definition's
|
||||
credentials somewhere else.
|
||||
"""
|
||||
entries: Final = stored_callback_entries(logging_entries)
|
||||
backends: Final = tuple(None if name is None else callback_backend_name(name) for name, _ in entries)
|
||||
return next((error for error in _reference_errors(entries, backends) if error is not None), None)
|
||||
|
||||
|
||||
def _reference_errors(
|
||||
entries: Sequence[tuple[str | None, Mapping[str, str]]], backends: Sequence[str | None]
|
||||
) -> Iterator[str | None]:
|
||||
for (name, callback_vars), backend in zip(entries, backends, strict=True):
|
||||
if name is None or callback_definition(name) is None:
|
||||
continue
|
||||
if callback_vars:
|
||||
yield f"callback_settings.{name} is referenced by name only, so the entry may not carry callback_vars"
|
||||
elif backends.count(backend) > 1:
|
||||
yield (
|
||||
f"callback_settings.{name} already provides the {backend} callback for this key, "
|
||||
f"so no other {backend} entry may sit next to it"
|
||||
)
|
||||
|
||||
|
||||
def logging_entry_keys(logging_entries: object) -> tuple[LoggingEntryKey, ...]:
|
||||
"""Every ``logging`` entry as a comparable value: name, event type and sorted vars, so no field changes unnoticed."""
|
||||
entries: Final = _object_sequence(logging_entries) or ()
|
||||
return tuple(
|
||||
(_entry_callback_name(entry), _entry_callback_type(entry), tuple(sorted(_entry_callback_vars(entry).items())))
|
||||
for entry in entries
|
||||
)
|
||||
|
||||
|
||||
def callback_config_error(
|
||||
|
|
@ -279,6 +425,9 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
|
|||
entries: Final = _object_sequence(metadata.get("logging"))
|
||||
if entries is None:
|
||||
return None
|
||||
reference_error: Final = callback_reference_error(entries)
|
||||
if reference_error is not None:
|
||||
return reference_error
|
||||
entry_error: Final = next(
|
||||
(error for entry in entries if (error := _logging_entry_error(entry)) is not None),
|
||||
None,
|
||||
|
|
@ -318,6 +467,14 @@ def _entry_callback_name(entry: object) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _entry_callback_type(entry: object) -> str | None:
|
||||
match entry:
|
||||
case {"callback_type": str() as callback_type}:
|
||||
return callback_type
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _conflicting_entry_setting_error(
|
||||
entry: tuple[str | None, Mapping[str, str]], earlier: Sequence[tuple[str | None, Mapping[str, str]]]
|
||||
) -> str | None:
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_utils import get_request_route
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.callback_config_validation import callback_backend_name, callback_definition
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
decrypt_callback_vars,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -1243,6 +1244,69 @@ def get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallba
|
|||
|
||||
_get_validated_callback_metadata: Final = get_validated_callback_metadata
|
||||
|
||||
_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def resolved_key_callbacks(logging_entries: object, *, source: str) -> tuple[AddTeamCallback, ...]:
|
||||
"""A key's ``logging`` entries as the backends see them, each ``callback_settings`` reference swapped for its definition.
|
||||
|
||||
A reference is honoured only while no other entry on the key names its backend: every
|
||||
entry's vars merge into one map per backend, so a sibling could point the definition's
|
||||
credentials at its own host. The reference's own ``callback_vars`` are never read.
|
||||
"""
|
||||
if not isinstance(logging_entries, Sequence) or isinstance(logging_entries, (str, bytes)):
|
||||
return ()
|
||||
entries: Final = tuple(
|
||||
_STRING_KEYED_MAPPING.validate_python(item)
|
||||
for item in _OBJECT_LIST_ADAPTER.validate_python(logging_entries)
|
||||
if isinstance(item, Mapping)
|
||||
)
|
||||
backends: Final = tuple(_entry_backend(entry) for entry in entries)
|
||||
return tuple(
|
||||
callback
|
||||
for entry, backend in zip(entries, backends, strict=True)
|
||||
if (callback := _resolved_key_callback(entry, source=source, shared_backend=backends.count(backend) > 1))
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
def _entry_backend(entry: Mapping[str, object]) -> str | None:
|
||||
callback_name: Final = entry.get("callback_name")
|
||||
return callback_backend_name(callback_name) if isinstance(callback_name, str) else None
|
||||
|
||||
|
||||
def _resolved_key_callback(entry: Mapping[str, object], *, source: str, shared_backend: bool) -> AddTeamCallback | None:
|
||||
definition: Final = callback_definition(entry.get("callback_name"))
|
||||
if definition is None:
|
||||
return get_validated_callback_metadata(item=dict(entry), source=source)
|
||||
if shared_backend:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring %s callback metadata naming callback_settings.%s: another entry on this key names the %s callback",
|
||||
source,
|
||||
_sanitize_for_log(definition.name),
|
||||
_sanitize_for_log(definition.backend),
|
||||
)
|
||||
return None
|
||||
reference: Final = get_validated_callback_metadata(item={**entry, "callback_vars": {}}, source=source)
|
||||
if reference is None:
|
||||
return None
|
||||
return reference.model_copy(
|
||||
update={"callback_name": definition.backend, "callback_vars": dict(definition.callback_vars)}
|
||||
)
|
||||
|
||||
|
||||
def _validated_callbacks(logging_entries: Iterable[object], *, source: str) -> tuple[AddTeamCallback, ...]:
|
||||
return tuple(
|
||||
callback
|
||||
for item in logging_entries
|
||||
if (
|
||||
callback := get_validated_callback_metadata(
|
||||
item=dict(_STRING_KEYED_MAPPING.validate_python(item)), source=source
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
class KeyAndTeamLoggingSettings:
|
||||
"""
|
||||
|
|
@ -1276,10 +1340,7 @@ def get_dynamic_logging_metadata(
|
|||
# Key-based callbacks
|
||||
#########################################################################################
|
||||
if key_dynamic_logging_settings is not None:
|
||||
for item in key_dynamic_logging_settings:
|
||||
callback = get_validated_callback_metadata(item=item, source="key-level")
|
||||
if callback is None:
|
||||
continue
|
||||
for callback in resolved_key_callbacks(key_dynamic_logging_settings, source="key-level"):
|
||||
callback_settings_obj = convert_key_logging_metadata_to_callback(
|
||||
data=callback,
|
||||
team_callback_settings_obj=callback_settings_obj,
|
||||
|
|
@ -1443,8 +1504,11 @@ def resolve_tenant_otel_destinations(
|
|||
disabled: Final = _dynamically_disabled_backends(user_api_key_dict, request_headers)
|
||||
callbacks: Final = tuple(
|
||||
callback
|
||||
for item in entries
|
||||
if (callback := get_validated_callback_metadata(item=item, source="otel-destination")) is not None
|
||||
for callback in (
|
||||
resolved_key_callbacks(entries, source="otel-destination")
|
||||
if key_entries is not None
|
||||
else _validated_callbacks(entries, source="otel-destination")
|
||||
)
|
||||
if callback.callback_name.lower() not in disabled
|
||||
)
|
||||
return tuple(
|
||||
|
|
|
|||
|
|
@ -76,6 +76,8 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq
|
|||
check_permissions_caller_permission,
|
||||
generate_key_helper_fn,
|
||||
prepare_metadata_fields,
|
||||
raise_on_invalid_key_logging_config,
|
||||
raise_on_unauthorized_callback_reference,
|
||||
)
|
||||
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
|
||||
|
|
@ -594,6 +596,11 @@ 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")
|
||||
raise_on_invalid_key_logging_config(data.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
raise_on_unauthorized_callback_reference(
|
||||
data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
user_api_key_dict,
|
||||
)
|
||||
|
||||
check_permissions_caller_permission(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -67,7 +67,11 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
|||
evict_and_broadcast,
|
||||
publish_auth_cache_invalidation,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
|
||||
from litellm.proxy.common_utils.callback_config_validation import (
|
||||
callback_definition,
|
||||
logging_entry_keys,
|
||||
logging_metadata_config_error,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
decrypt_callback_vars,
|
||||
encrypt_callback_vars,
|
||||
|
|
@ -433,6 +437,7 @@ _KEY_METADATA_REQUEST_FIELDS: Final = frozenset(
|
|||
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_STRING_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _decode_json_string_column(column: str, value: object) -> object:
|
||||
|
|
@ -933,12 +938,49 @@ def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -
|
|||
raise HTTPException(status_code=400, detail={"error": error})
|
||||
|
||||
|
||||
def raise_on_unauthorized_callback_reference(
|
||||
metadata: Mapping[str, object] | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
*,
|
||||
existing_metadata: Mapping[str, object] | None = None,
|
||||
) -> None:
|
||||
"""Only a proxy admin may point a key at a ``callback_settings`` definition or change the logging of a key holding one.
|
||||
|
||||
The definition carries the operator's credentials, so the team admins and members who
|
||||
can otherwise edit key metadata may only send such a key's ``logging`` list back unchanged.
|
||||
"""
|
||||
if metadata is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
requested: Final = logging_entry_keys(metadata.get("logging"))
|
||||
stored: Final = logging_entry_keys(_stored_logging_entries(existing_metadata))
|
||||
if requested == stored or not any(callback_definition(name) is not None for name, _, _ in (*requested, *stored)):
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"Only a proxy admin can point a key at a callback_settings definition "
|
||||
"or change the logging entries of a key that holds one"
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _stored_logging_entries(existing_metadata: Mapping[str, object] | None) -> object:
|
||||
"""The stored ``logging`` list with its callback vars decrypted, since ciphertext never equals the plaintext sent back."""
|
||||
if existing_metadata is None:
|
||||
return None
|
||||
decrypted: Final = _DECODED_JSON.validate_python(decrypt_callback_vars(existing_metadata))
|
||||
return _STRING_KEYED_MAPPING.validate_python(decrypted).get("logging") if isinstance(decrypted, Mapping) else None
|
||||
|
||||
|
||||
def common_key_access_checks(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
llm_router: Router | None,
|
||||
premium_user: bool,
|
||||
user_id: str | None = None,
|
||||
existing_metadata: Mapping[str, object] | None = None,
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
Check if user is allowed to make a key request, for this key
|
||||
|
|
@ -965,6 +1007,11 @@ def common_key_access_checks(
|
|||
llm_router=llm_router,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
raise_on_unauthorized_callback_reference(
|
||||
data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
user_api_key_dict,
|
||||
existing_metadata=existing_metadata,
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
|
|
@ -3773,6 +3820,7 @@ async def _validate_update_key_data(
|
|||
user_id=user_api_key_dict.user_id if acting_as_team_admin else existing_key_row.user_id,
|
||||
llm_router=llm_router,
|
||||
premium_user=premium_user,
|
||||
existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
)
|
||||
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
|
|
@ -6652,6 +6700,11 @@ async def regenerate_key_fn(
|
|||
_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
user_api_key_dict,
|
||||
)
|
||||
raise_on_unauthorized_callback_reference(
|
||||
None if data is None else data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict
|
||||
user_api_key_dict,
|
||||
existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict
|
||||
)
|
||||
|
||||
# check if user has permission to regenerate key
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
|
|
|
|||
|
|
@ -168,6 +168,7 @@ from litellm.proxy._types import (
|
|||
TransformRequestBody,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_config_validation import callback_definitions_error
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
is_sensitive_callback_key,
|
||||
normalize_callback_names,
|
||||
|
|
@ -5862,6 +5863,16 @@ def _configured_db_callbacks(litellm_settings: Mapping[str, object], setting_key
|
|||
return tuple((setting_key, callback) for callback in callbacks if isinstance(callback, str))
|
||||
|
||||
|
||||
_CALLBACK_SETTINGS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _check_callback_definitions(callback_settings: Mapping[str, object]) -> None:
|
||||
"""Refuse to boot on a ``callback_settings`` definition a key could never use."""
|
||||
definition_error: Final = callback_definitions_error(callback_settings)
|
||||
if definition_error is not None:
|
||||
raise ValueError(definition_error)
|
||||
|
||||
|
||||
class ProxyConfig:
|
||||
"""
|
||||
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
||||
|
|
@ -6785,6 +6796,7 @@ class ProxyConfig:
|
|||
## Callback settings
|
||||
callback_settings: Final = config.get("callback_settings", {})
|
||||
if callback_settings:
|
||||
_check_callback_definitions(_CALLBACK_SETTINGS.validate_python(callback_settings))
|
||||
litellm.callback_settings = callback_settings
|
||||
|
||||
## LITELLM MODULE SETTINGS (e.g. litellm.drop_params=True,..)
|
||||
|
|
|
|||
|
|
@ -3,11 +3,16 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.common_utils.callback_config_validation import (
|
||||
callback_backend_name,
|
||||
callback_config_error,
|
||||
callback_definition,
|
||||
callback_definitions_error,
|
||||
conflicting_callback_var_error,
|
||||
conflicting_span_scope_error,
|
||||
cross_entry_family_error,
|
||||
logging_entry_keys,
|
||||
logging_metadata_config_error,
|
||||
)
|
||||
|
||||
|
|
@ -499,3 +504,158 @@ 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
|
||||
|
||||
|
||||
_ARIZE_DEFINITION: Final = {
|
||||
"callback_type": "arize",
|
||||
"arize_api_key": "resolved-space-a-key",
|
||||
"arize_space_id": "space-a",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"arize_space_a": dict(_ARIZE_DEFINITION)})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("callback_settings", "named_fields"),
|
||||
[
|
||||
({"arize": {"callback_type": "arize", "arize_api_key": "k"}}, ("callback_settings.arize", "named after")),
|
||||
(
|
||||
{"space_a": {"callback_type": "not_a_callback"}},
|
||||
("callback_settings.space_a.callback_type", "not_a_callback"),
|
||||
),
|
||||
({"space_a": {"callback_type": "arize", "arize_api_key": None}}, ("callback_settings.space_a.arize_api_key",)),
|
||||
({"space_a": {"callback_type": "arize", "arize_api_key": ""}}, ("callback_settings.space_a.arize_api_key",)),
|
||||
(
|
||||
{"space_a": {"callback_type": "arize", "arize_space_id": "os.environ/UNSET"}},
|
||||
("callback_settings.space_a.arize_space_id", "secret manager"),
|
||||
),
|
||||
({"space_a": {"callback_type": "arize", "arize_api_kye": "k"}}, ("callback_settings.space_a", "arize_api_kye")),
|
||||
(
|
||||
{"space_a": {"callback_type": "arize", "arize_otlp_protocol": "carrier-pigeon"}},
|
||||
("callback_settings.space_a", "arize_otlp_protocol"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_definition_a_key_could_never_use_fails_config_load(
|
||||
callback_settings: dict[str, object], named_fields: tuple[str, ...]
|
||||
) -> None:
|
||||
error: Final = callback_definitions_error(callback_settings)
|
||||
assert error is not None
|
||||
for fragment in named_fields:
|
||||
assert fragment in error
|
||||
|
||||
|
||||
def test_a_definition_error_names_the_field_and_never_its_value() -> None:
|
||||
error: Final = callback_definitions_error(
|
||||
{
|
||||
"space_a": {
|
||||
"callback_type": "arize",
|
||||
"arize_otlp_protocol": "carrier-pigeon",
|
||||
"arize_api_key": "sk-super-secret",
|
||||
}
|
||||
}
|
||||
)
|
||||
assert error is not None and "arize_otlp_protocol" in error
|
||||
assert "sk-super-secret" not in error
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("otel_v2_on")
|
||||
def test_callback_settings_without_a_referenceable_definition_pass_config_load() -> None:
|
||||
assert (
|
||||
callback_definitions_error(
|
||||
{
|
||||
"otel": {"message_logging": False, "excluded_services": ["postgres"]},
|
||||
"my_api": {"callback_type": "generic_api", "url": "os.environ/UNSET"},
|
||||
"space_a": {
|
||||
"callback_type": "arize",
|
||||
"arize_api_key": "k",
|
||||
"arize_space_id": 12345,
|
||||
"excluded_services": "postgres,redis",
|
||||
},
|
||||
}
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_reference_resolves_to_its_definition_and_a_plain_name_does_not() -> None:
|
||||
definition: Final = callback_definition("arize_space_a")
|
||||
assert definition is not None
|
||||
assert (definition.backend, dict(definition.callback_vars)) == (
|
||||
"arize",
|
||||
{"arize_api_key": "resolved-space-a-key", "arize_space_id": "space-a"},
|
||||
)
|
||||
assert callback_definition("arize") is None
|
||||
assert callback_definition(None) is None
|
||||
assert (callback_backend_name("arize_space_a"), callback_backend_name("langfuse")) == ("arize", "langfuse")
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_reference_entry_may_not_carry_callback_vars() -> None:
|
||||
error: Final = logging_metadata_config_error(
|
||||
{
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "arize_space_a",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"arize_space_id": "attacker-space"},
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert error is not None and "arize_space_a" in error and "callback_vars" in error
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_reference_is_the_only_entry_for_its_backend() -> None:
|
||||
sibling: Final = {"callback_name": "arize", "callback_type": "success", "callback_vars": {"arize_space_id": "x"}}
|
||||
reference: Final = {"callback_name": "arize_space_a", "callback_type": "success"}
|
||||
error: Final = logging_metadata_config_error({"logging": [reference, sibling]})
|
||||
assert error is not None and "arize_space_a" in error and "arize" in error
|
||||
assert logging_metadata_config_error({"logging": [sibling]}) is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_reference_sits_next_to_entries_for_other_backends() -> None:
|
||||
assert (
|
||||
logging_metadata_config_error(
|
||||
{
|
||||
"logging": [
|
||||
{"callback_name": "arize_space_a", "callback_type": "success"},
|
||||
{"callback_name": "arize_space_a", "callback_type": "failure", "callback_vars": {}},
|
||||
{"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}},
|
||||
]
|
||||
}
|
||||
)
|
||||
is not None
|
||||
), "two references to one definition still share its backend"
|
||||
assert (
|
||||
logging_metadata_config_error(
|
||||
{
|
||||
"logging": [
|
||||
{"callback_name": "arize_space_a", "callback_type": "success"},
|
||||
{"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}},
|
||||
]
|
||||
}
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_logging_entry_keys_change_with_every_field_a_caller_could_edit() -> None:
|
||||
entry: Final = {"callback_name": "arize_space_a", "callback_type": "success", "callback_vars": {"b": "2", "a": "1"}}
|
||||
reordered: Final = {
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"a": "1", "b": "2"},
|
||||
"callback_name": "arize_space_a",
|
||||
}
|
||||
assert logging_entry_keys([entry]) == logging_entry_keys([reordered])
|
||||
assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_type": "failure"}])
|
||||
assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_vars": {"a": "1"}}])
|
||||
assert logging_entry_keys([entry]) != logging_entry_keys([{**entry, "callback_name": "arize"}])
|
||||
assert logging_entry_keys([entry]) != logging_entry_keys([entry, entry])
|
||||
assert logging_entry_keys(None) == logging_entry_keys("not-a-list") == ()
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ from fastapi.testclient import TestClient
|
|||
from pydantic import ValidationError
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_UserTable,
|
||||
|
|
@ -5496,3 +5498,42 @@ async def test_fips_user_update_hashes_off_the_event_loop(pbkdf2_on_event_loop:
|
|||
|
||||
assert pbkdf2_on_event_loop == [False]
|
||||
assert data["password"].startswith("pbkdf2:")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_user_cannot_hand_its_auto_created_key_a_callback_reference_for_a_non_admin(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""``/user/new`` writes ``metadata`` onto the key it auto-creates, so it runs the same admin gate as ``/key/generate``."""
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callback_settings",
|
||||
{"arize_space_a": {"callback_type": "arize", "arize_api_key": "k", "arize_space_id": "s"}},
|
||||
)
|
||||
table: Final = SimpleNamespace(
|
||||
find_first=AsyncMock(return_value=None),
|
||||
count=AsyncMock(return_value=0),
|
||||
create=AsyncMock(),
|
||||
upsert=AsyncMock(),
|
||||
)
|
||||
mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table))
|
||||
database: Final = PrismaClient(
|
||||
database_url="postgresql://unused:unused@localhost/unused",
|
||||
proxy_logging_obj=proxy_server.proxy_logging_obj,
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", database)
|
||||
reference: Final = {"logging": [{"callback_name": "arize_space_a", "callback_type": "success"}]}
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await new_user(
|
||||
data=NewUserRequest(user_email="bob@example.test", metadata=reference, send_invite_email=False),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "callback_settings" in exc.value.message
|
||||
table.create.assert_not_awaited()
|
||||
table.upsert.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
|
|
@ -87,6 +88,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
_validate_max_budget,
|
||||
_validate_reset_spend_value,
|
||||
_validate_update_key_data,
|
||||
raise_on_unauthorized_callback_reference,
|
||||
can_modify_verification_token,
|
||||
check_org_key_model_specific_limits,
|
||||
check_team_key_model_specific_limits,
|
||||
|
|
@ -24243,3 +24245,187 @@ 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")
|
||||
|
||||
|
||||
_REFERENCE_ENTRY: Final = {"callback_name": "arize_space_a", "callback_type": "success"}
|
||||
_LANGFUSE_ENTRY: Final = {
|
||||
"callback_name": "langfuse",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk-plain"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callback_settings",
|
||||
{"arize_space_a": {"callback_type": "arize", "arize_api_key": "k", "arize_space_id": "s"}},
|
||||
)
|
||||
|
||||
|
||||
def _non_admin() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-alice", user_id="alice")
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
class TestCallbackReferenceWriteGate:
|
||||
def test_a_proxy_admin_may_point_a_key_at_a_definition(self) -> None:
|
||||
admin: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
assert raise_on_unauthorized_callback_reference({"logging": [_REFERENCE_ENTRY]}, admin) is None
|
||||
assert (
|
||||
raise_on_unauthorized_callback_reference(
|
||||
{"logging": []}, admin, existing_metadata={"logging": [_REFERENCE_ENTRY]}
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata",
|
||||
[
|
||||
{"logging": [_REFERENCE_ENTRY]},
|
||||
{"logging": [_LANGFUSE_ENTRY, _REFERENCE_ENTRY]},
|
||||
{"logging": [{**_REFERENCE_ENTRY, "callback_vars": {}}]},
|
||||
],
|
||||
)
|
||||
def test_a_non_admin_cannot_point_a_key_at_a_definition(self, metadata: Mapping[str, object]) -> None:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_unauthorized_callback_reference(metadata, _non_admin())
|
||||
assert exc.value.status_code == 403
|
||||
assert "callback_settings" in str(exc.value.detail)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata",
|
||||
[
|
||||
{"note": "edited"},
|
||||
{"logging": []},
|
||||
{"logging": [{**_REFERENCE_ENTRY, "callback_type": "failure"}]},
|
||||
{"logging": [{**_REFERENCE_ENTRY, "callback_name": "arize"}]},
|
||||
{"logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]},
|
||||
{"logging": [{**_LANGFUSE_ENTRY, "callback_vars": {"langfuse_host": "https://attacker.example"}}]},
|
||||
],
|
||||
)
|
||||
def test_a_non_admin_cannot_change_the_logging_of_a_key_holding_a_reference(
|
||||
self, metadata: Mapping[str, object]
|
||||
) -> None:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
raise_on_unauthorized_callback_reference(
|
||||
metadata, _non_admin(), existing_metadata={"logging": [_REFERENCE_ENTRY]}
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
|
||||
def test_a_non_admin_may_edit_the_rest_of_the_metadata_when_the_logging_is_sent_back_unchanged(self) -> None:
|
||||
stored: Final = encrypt_callback_vars({"note": "original", "logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]})
|
||||
assert (
|
||||
raise_on_unauthorized_callback_reference(
|
||||
{"note": "edited", "logging": [_REFERENCE_ENTRY, _LANGFUSE_ENTRY]},
|
||||
_non_admin(),
|
||||
existing_metadata=stored,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("existing_metadata", [None, {}, {"logging": [_LANGFUSE_ENTRY]}])
|
||||
def test_a_key_without_a_reference_keeps_the_logging_rules_it_had(
|
||||
self, existing_metadata: Mapping[str, object] | None
|
||||
) -> None:
|
||||
changed: Final = [
|
||||
{**_LANGFUSE_ENTRY, "callback_type": "failure"},
|
||||
{"callback_name": "datadog", "callback_vars": {}},
|
||||
]
|
||||
assert (
|
||||
raise_on_unauthorized_callback_reference(
|
||||
{"logging": changed}, _non_admin(), existing_metadata=existing_metadata
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
raise_on_unauthorized_callback_reference(
|
||||
None, _non_admin(), existing_metadata={"logging": [_REFERENCE_ENTRY]}
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
async def test_generate_key_with_a_callback_reference_is_admin_only() -> None:
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(metadata={"logging": [_REFERENCE_ENTRY]}),
|
||||
user_api_key_dict=_non_admin(),
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "callback_settings" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
async def test_generate_key_rejects_callback_vars_on_a_reference_for_everyone() -> None:
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(
|
||||
metadata={"logging": [{**_REFERENCE_ENTRY, "callback_vars": {"arize_space_id": "mine"}}]}
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
assert exc.value.status_code == 400
|
||||
assert "callback_vars" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
async def test_update_key_compares_the_sent_logging_against_the_stored_row() -> None:
|
||||
existing_key_row: Final = MagicMock()
|
||||
existing_key_row.user_id = "alice"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
existing_key_row.metadata = {"logging": [_REFERENCE_ENTRY]}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _validate_update_key_data(
|
||||
data=UpdateKeyRequest(key="sk-test-key", metadata={"note": "edited"}),
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=_non_admin(),
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "callback_settings" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
async def test_regenerate_key_cannot_drop_a_reference_for_a_non_admin(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import regenerate_key_fn
|
||||
|
||||
key_row: Final = MagicMock()
|
||||
key_row.token = "hashed_token"
|
||||
key_row.user_id = "alice"
|
||||
key_row.team_id = None
|
||||
key_row.metadata = {"logging": [_REFERENCE_ENTRY]}
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row)
|
||||
prisma.db.litellm_verificationtoken.update = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master")
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await regenerate_key_fn(
|
||||
key=None,
|
||||
data=RegenerateKeyRequest(key="sk-test", metadata={"note": "edited"}),
|
||||
user_api_key_dict=_non_admin(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "callback_settings" in str(exc.value.message)
|
||||
prisma.db.litellm_verificationtoken.update.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -6593,3 +6593,68 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path,
|
|||
return
|
||||
with pytest.raises(ValidationError):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config))
|
||||
|
||||
|
||||
_CALLBACK_DEFINITION_CONFIG: Final = """
|
||||
model_list: []
|
||||
callback_settings:
|
||||
arize_space_a:
|
||||
callback_type: arize
|
||||
arize_api_key: os.environ/QA_ARIZE_SPACE_A_KEY
|
||||
arize_space_id: os.environ/QA_ARIZE_SPACE_A_SPACE
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config_without_db(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
monkeypatch.setattr(litellm, "callback_settings", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("config_without_db")
|
||||
async def test_a_callback_definition_resolves_its_secrets_at_boot(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("QA_ARIZE_SPACE_A_KEY", "space-a-key-v1")
|
||||
monkeypatch.setenv("QA_ARIZE_SPACE_A_SPACE", "space-a")
|
||||
config_file: Final = tmp_path / "callbacks.yaml"
|
||||
config_file.write_text(_CALLBACK_DEFINITION_CONFIG)
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config_file))
|
||||
|
||||
assert litellm.callback_settings["arize_space_a"] == {
|
||||
"callback_type": "arize",
|
||||
"arize_api_key": "space-a-key-v1",
|
||||
"arize_space_id": "space-a",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("config_without_db")
|
||||
async def test_a_callback_definition_whose_secret_is_missing_fails_boot_naming_the_field(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("QA_ARIZE_SPACE_A_KEY", "space-a-key-v1")
|
||||
monkeypatch.delenv("QA_ARIZE_SPACE_A_SPACE", raising=False)
|
||||
config_file: Final = tmp_path / "callbacks.yaml"
|
||||
config_file.write_text(_CALLBACK_DEFINITION_CONFIG)
|
||||
|
||||
with pytest.raises(ValueError, match=r"callback_settings\.arize_space_a\.arize_space_id"):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config_file))
|
||||
|
||||
assert litellm.callback_settings == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("config_without_db")
|
||||
async def test_a_callback_definition_with_a_misspelled_field_fails_boot(tmp_path: Path) -> None:
|
||||
config_file: Final = tmp_path / "callbacks.yaml"
|
||||
config_file.write_text(
|
||||
"model_list: []\ncallback_settings:\n arize_space_a:\n callback_type: arize\n arize_api_kye: k\n"
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="arize_api_kye"):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config_file))
|
||||
|
|
|
|||
|
|
@ -11443,3 +11443,111 @@ def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json
|
|||
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data)
|
||||
assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"}
|
||||
assert data["metadata"] == caller_metadata
|
||||
|
||||
|
||||
_REFERENCE_ENTRY: Final = {"callback_name": "arize_space_a", "callback_type": "success"}
|
||||
_ARIZE_SPACE_A: Final = {
|
||||
"callback_type": "arize",
|
||||
"arize_api_key": "resolved-space-a-key",
|
||||
"arize_space_id": "space-a",
|
||||
"arize_otlp_protocol": "http/protobuf",
|
||||
"excluded_services": "postgres,redis",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def arize_space_a(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "callback_settings", {"arize_space_a": dict(_ARIZE_SPACE_A)})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def otel_v2_on(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setenv("ARIZE_ENDPOINT", "https://arize.internal.example/v1")
|
||||
monkeypatch.delenv("ARIZE_HTTP_ENDPOINT", raising=False)
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
yield
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_key_logging_entry_naming_a_definition_runs_its_backend_with_the_definition_vars() -> None:
|
||||
auth: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={"logging": [_REFERENCE_ENTRY]}, team_metadata={})
|
||||
|
||||
result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock())
|
||||
|
||||
assert result is not None
|
||||
assert result.success_callback == ["arize"]
|
||||
assert not result.failure_callback
|
||||
assert result.callback_vars == {
|
||||
key: value for key, value in _ARIZE_SPACE_A.items() if key not in ("callback_type", "excluded_services")
|
||||
}, "excluded_services is a destination-only var, which the destination test below covers"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_the_vars_stored_on_a_reference_entry_never_reach_the_backend() -> None:
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
metadata={"logging": [{**_REFERENCE_ENTRY, "callback_vars": {"arize_space_id": "attacker-space"}}]},
|
||||
team_metadata={},
|
||||
)
|
||||
|
||||
result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock())
|
||||
|
||||
assert result is not None and result.callback_vars is not None
|
||||
assert result.callback_vars["arize_space_id"] == "space-a"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_reference_is_skipped_when_another_entry_on_the_key_names_its_backend() -> None:
|
||||
"""A stored inline entry for the same backend would merge its host or space with the definition's
|
||||
credentials, so the definition's secrets stay out of the request instead."""
|
||||
sibling: Final = {
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {"arize_space_id": "attacker-space", "arize_api_key": "attacker-key"},
|
||||
}
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-key", metadata={"logging": [_REFERENCE_ENTRY, sibling]}, team_metadata={}
|
||||
)
|
||||
|
||||
result: Final = get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock())
|
||||
|
||||
assert result is not None
|
||||
assert result.callback_vars == sibling["callback_vars"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a")
|
||||
def test_a_team_logging_entry_naming_a_definition_is_not_resolved() -> None:
|
||||
auth: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={"logging": [_REFERENCE_ENTRY]})
|
||||
|
||||
assert get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock()) is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("arize_space_a", "otel_v2_on")
|
||||
def test_a_reference_resolves_to_the_same_otel_destination_as_the_inline_entry() -> None:
|
||||
from litellm.proxy.litellm_pre_call_utils import resolve_tenant_otel_destinations
|
||||
|
||||
inline: Final = {
|
||||
"callback_name": "arize",
|
||||
"callback_type": "success",
|
||||
"callback_vars": {key: value for key, value in _ARIZE_SPACE_A.items() if key != "callback_type"},
|
||||
}
|
||||
referencing_key: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-key", team_id="team-a", metadata={"logging": [_REFERENCE_ENTRY]}, team_metadata={}
|
||||
)
|
||||
inline_key: Final = UserAPIKeyAuth(
|
||||
api_key="other-key", team_id="team-a", metadata={"logging": [inline]}, team_metadata={}
|
||||
)
|
||||
plain_key_on_the_same_team: Final = UserAPIKeyAuth(
|
||||
api_key="plain-key", team_id="team-a", metadata={}, team_metadata={}
|
||||
)
|
||||
|
||||
resolved: Final = resolve_tenant_otel_destinations(referencing_key)
|
||||
|
||||
assert resolved == resolve_tenant_otel_destinations(inline_key)
|
||||
assert len(resolved) == 1
|
||||
assert resolved[0].headers["api_key"] == "resolved-space-a-key"
|
||||
assert resolve_tenant_otel_destinations(plain_key_on_the_same_team) == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue