mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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
This commit is contained in:
Jesus Velarde
parent
d0c16008c8
commit
b4a7baec27
12 changed files with 554 additions and 25 deletions
|
|
@ -2099,6 +2099,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"
|
||||
|
|
@ -2108,14 +2113,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,71 @@ 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)
|
||||
references_involved: Final = bool(requested_names or logging_credential_names(stored_logging))
|
||||
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 requested_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,44 @@ 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 object_sequence(entries) or ()
|
||||
if (resolved := _with_credential(string_keyed(entry))) is not None
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
if CredentialAccessor.find_credential(credential_name) is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring key logging entry for %s: credential %r is not defined",
|
||||
entry.get("callback_name"),
|
||||
credential_name,
|
||||
)
|
||||
return None
|
||||
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
|
||||
}
|
||||
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
|
||||
|
|
@ -1237,9 +1238,9 @@ 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: Mapping[str, object], *, source: str) -> AddTeamCallback | None:
|
||||
try:
|
||||
return AddTeamCallback(**item)
|
||||
return AddTeamCallback.model_validate(item)
|
||||
except (PydanticValidationError, ValueError) as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid %s callback metadata: %s",
|
||||
|
|
@ -1258,9 +1259,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
|
||||
|
|
@ -1274,9 +1275,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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -76,6 +76,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
|
||||
|
|
@ -594,6 +595,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")
|
||||
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
|
||||
|
|
@ -67,9 +67,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 (
|
||||
|
|
@ -264,6 +269,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]
|
||||
|
||||
|
|
@ -932,6 +943,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,
|
||||
|
|
@ -1483,6 +1511,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(
|
||||
|
|
@ -3414,6 +3443,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
|
||||
)
|
||||
|
|
@ -4153,6 +4183,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
|
||||
)
|
||||
|
|
@ -6327,6 +6358,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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,222 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from itertools import accumulate
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from _openinference_support import (
|
||||
_chat_request_marker,
|
||||
_chat_response,
|
||||
_matching_marker_spans,
|
||||
_rig,
|
||||
)
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, Wire
|
||||
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": {},
|
||||
}
|
||||
]
|
||||
}
|
||||
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 _chat(proxy: Gateway, model: str, key: str, marker: str) -> None:
|
||||
response: Final = proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"metadata": {"trace_marker": marker},
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
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
|
||||
|
|
@ -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,44 @@ 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
|
||||
|
|
|
|||
|
|
@ -21,6 +21,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 +29,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 +777,43 @@ 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"},
|
||||
]
|
||||
|
||||
|
||||
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,)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -24243,3 +24246,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
|
||||
|
|
|
|||
|
|
@ -17,8 +17,6 @@ from pydantic import ValidationError as PydanticValidationError
|
|||
from starlette.datastructures import Headers
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY,
|
||||
|
|
@ -28,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,
|
||||
|
|
@ -40,6 +39,7 @@ from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_log
|
|||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
|
||||
from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
KeyAndTeamLoggingSettings,
|
||||
LiteLLMProxyRequestSetup,
|
||||
|
|
@ -60,10 +60,12 @@ 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
|
||||
from litellm.types.utils import CredentialItem
|
||||
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
||||
|
||||
|
||||
def test_check_if_token_is_service_account():
|
||||
|
|
@ -2776,7 +2778,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(
|
||||
|
|
@ -2785,7 +2787,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={})
|
||||
|
|
@ -10194,6 +10196,49 @@ 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