Skip to content
Merged
16 changes: 15 additions & 1 deletion litellm/litellm_core_utils/litellm_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,10 @@
_PAGERDUTY_ALERTING_FACTORY: Final = PagerDutyAlerting
_in_memory_loggers: Final[list[CustomLogger]] = []

_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys())
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(("used_client_oauth_token",))
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = (
frozenset(StandardLoggingMetadata.__annotations__.keys()) - _STANDARD_LOGGING_METADATA_RESOLVED_KEYS
)


def _get_provider_request_id(original_exception: Exception) -> str | None:
Expand Down Expand Up @@ -5715,6 +5718,7 @@ def get_standard_logging_metadata(
proxy_server_request: dict | None = None,
start_time: dt_object | None = None,
response_id: str | None = None,
custom_llm_provider: str | None = None,
) -> StandardLoggingMetadata:
"""
Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata.
Expand All @@ -5729,6 +5733,9 @@ def get_standard_logging_metadata(
- If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned.
- If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'.
"""
from litellm.llms.anthropic.common_utils import ( # noqa: PLC0415 # that module imports this one transitively
resolve_used_client_oauth_token,
)

prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None
if litellm_params is not None:
Expand Down Expand Up @@ -5778,6 +5785,10 @@ def get_standard_logging_metadata(
user_api_key_auth_metadata=None,
team_alias=None,
team_id=None,
used_client_oauth_token=resolve_used_client_oauth_token(
metadata.get("used_client_oauth_token") if isinstance(metadata, dict) else None,
custom_llm_provider,
),
Comment thread
cursor[bot] marked this conversation as resolved.
Resolved
)
if isinstance(metadata, dict):
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
Expand Down Expand Up @@ -6501,6 +6512,7 @@ def get_standard_logging_object_payload(
stream=kwargs.get("stream", False),
)
# clean up litellm metadata
selected_provider: Final = kwargs.get("custom_llm_provider")
clean_metadata: Final = StandardLoggingPayloadSetup.get_standard_logging_metadata(
metadata=metadata,
litellm_params=litellm_params,
Expand All @@ -6512,6 +6524,7 @@ def get_standard_logging_object_payload(
proxy_server_request=proxy_server_request,
start_time=start_time,
response_id=id,
custom_llm_provider=selected_provider if isinstance(selected_provider, str) else None,
)
_request_body: Final = proxy_server_request.get("body", {})
end_user_id: Final = clean_metadata["user_api_key_end_user_id"] or _request_body.get(
Expand Down Expand Up @@ -6786,6 +6799,7 @@ def get_standard_logging_metadata(
user_api_key_auth_metadata=None,
team_alias=None,
team_id=None,
used_client_oauth_token=None,
)
if isinstance(metadata, dict):
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields
Expand Down
10 changes: 10 additions & 0 deletions litellm/llms/anthropic/common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.model_listing import ModelInfoResponse
from litellm.types.utils import LlmProviders

_MessageT = TypeVar("_MessageT")

Expand Down Expand Up @@ -225,6 +226,15 @@ def is_anthropic_oauth_key(value: str | None) -> bool:
return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX)


ANTHROPIC_OAUTH_FORWARD_PROVIDERS: Final[frozenset[str]] = frozenset((LlmProviders.ANTHROPIC.value,))


def resolve_used_client_oauth_token(client_sent_oauth_token: object, custom_llm_provider: str | None) -> bool | None:
if not isinstance(client_sent_oauth_token, bool):
return None
return client_sent_oauth_token and custom_llm_provider in ANTHROPIC_OAUTH_FORWARD_PROVIDERS


def _merge_beta_headers(existing: str | None, new_beta: str) -> str:
"""Merge a new beta value into an existing comma-separated anthropic-beta header."""
if not existing:
Expand Down
1 change: 1 addition & 0 deletions litellm/proxy/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -4116,6 +4116,7 @@ class SpendLogsMetadata(TypedDict):
litellm_gateway_injected_cache: ReadOnly[str | None]
router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model
azure_spillover: ReadOnly[AzureSpillover | None] # None = Azure did not report spillover
used_client_oauth_token: ReadOnly[bool | None] # None = row written before the flag existed


class SpendLogsPayload(TypedDict):
Expand Down
9 changes: 9 additions & 0 deletions litellm/proxy/hooks/proxy_track_cost_callback.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
_get_parent_otel_span_from_kwargs,
budget_reservation_from_metadata,
get_litellm_metadata_from_kwargs,
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
Expand Down Expand Up @@ -85,6 +86,12 @@
)


def _proxy_stamped_used_client_oauth_token(request_data: Mapping[str, object]) -> bool | None:
proxy_metadata: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data))
stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None
return stamped if isinstance(stamped, bool) else None


def _proxy_spend_writer() -> DBSpendUpdateWriter:
from litellm.proxy.proxy_server import proxy_logging_obj

Expand Down Expand Up @@ -191,6 +198,8 @@ async def async_post_call_failure_hook(
metadata=_metadata, original_exception=original_exception
)

_metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_data)

existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
existing_metadata.update(_metadata)

Expand Down
10 changes: 7 additions & 3 deletions litellm/proxy/litellm_pre_call_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
is_url_destination_allowed_by_host,
provider_url_destination_candidates,
)
from litellm.llms.anthropic.common_utils import ANTHROPIC_OAUTH_FORWARD_PROVIDERS
from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
Expand Down Expand Up @@ -2175,7 +2176,9 @@ async def add_litellm_data_to_request(
data["api_version"] = dynamic_api_version

## Forward any LLM API Provider specific headers in extra_headers
add_provider_specific_headers_to_request(data=data, headers=_headers)
data[_metadata_variable_name]["used_client_oauth_token"] = add_provider_specific_headers_to_request(
data=data, headers=_headers
)
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved

## Cache Controls
cache_control_header: Final = _headers.get("Cache-Control", None)
Expand Down Expand Up @@ -3467,13 +3470,13 @@ async def add_guardrails_from_policy_engine(
LlmProviders.VERTEX_AI.value,
)
)
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.ANTHROPIC.value
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = ",".join(sorted(ANTHROPIC_OAUTH_FORWARD_PROVIDERS))
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved


def add_provider_specific_headers_to_request(
data: dict,
headers: dict,
):
) -> bool:
from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key

anthropic_api_headers: Final = {header: headers[header] for header in ANTHROPIC_API_HEADERS if header in headers}
Expand All @@ -3494,6 +3497,7 @@ def add_provider_specific_headers_to_request(

if scoped_headers:
data["provider_specific_header"] = scoped_headers[0] if len(scoped_headers) == 1 else scoped_headers
return bool(anthropic_oauth_credential_headers)


def _add_otel_traceparent_to_data(data: dict, request: Request):
Expand Down
13 changes: 13 additions & 0 deletions litellm/proxy/spend_tracking/spend_management_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -2515,6 +2515,15 @@ async def ui_view_spend_logs(
default=None,
description="Filter logs by cache state: 'hit' or 'miss'. Miss includes legacy rows with a null/unknown cache state",
),
used_client_oauth_token: Annotated[
bool | None,
fastapi.Query(
description=(
"Filter logs by the credential the upstream call used: true for a client-forwarded Anthropic OAuth "
"token, false for the deployment's configured key. Rows written before this flag existed match neither"
),
),
] = None,
span_type: str | None = fastapi.Query(
default=None,
description="Filter logs by span type: llm, agent, mcp, or batch",
Expand Down Expand Up @@ -2924,6 +2933,10 @@ def parse_date(date_str: str) -> datetime:
sql_conditions.append(f"metadata->'error_information'->>'error_message' LIKE ${p}")
sql_params.append(f"%{error_message}%")
p += 1
if used_client_oauth_token is not None:
sql_conditions.append(f"metadata->>'used_client_oauth_token' = ${p}")
sql_params.append(json.dumps(used_client_oauth_token))
p += 1

if (
group_by_session is True
Expand Down
8 changes: 8 additions & 0 deletions litellm/proxy/spend_tracking/spend_tracking_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
)
from litellm.litellm_core_utils.ptu_pricing import azure_spillover
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes
from litellm.llms.anthropic.common_utils import resolve_used_client_oauth_token
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
Expand Down Expand Up @@ -154,6 +155,7 @@ def _get_router_metadata_for_spend_log(
"autorouter_savings",
"autorouter_savings_estimate",
"autorouter_baseline_observation",
"used_client_oauth_token",
)
)

Expand All @@ -178,6 +180,7 @@ def _get_spend_logs_metadata(
autorouter_baseline_observation: str | None = None,
router_metadata: SpendLogsRouterMetadata | None = None,
azure_spillover: AzureSpillover | None = None,
used_client_oauth_token: bool | None = None,
) -> SpendLogsMetadata:
if metadata is None:
return SpendLogsMetadata(
Expand Down Expand Up @@ -222,6 +225,7 @@ def _get_spend_logs_metadata(
litellm_call_id=litellm_call_id,
router_metadata=router_metadata,
azure_spillover=azure_spillover,
used_client_oauth_token=used_client_oauth_token,
)
verbose_proxy_logger.debug(
"getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys()))
Expand All @@ -237,6 +241,7 @@ def _get_spend_logs_metadata(
autorouter_baseline_observation=autorouter_baseline_observation,
router_metadata=router_metadata,
azure_spillover=azure_spillover,
used_client_oauth_token=used_client_oauth_token,
)
_raw_key: Final = clean_metadata.get("user_api_key")
_trusted_hash: Final = metadata.get("user_api_key_hash")
Expand Down Expand Up @@ -714,6 +719,9 @@ def get_logging_payload(
selected_provider=custom_llm_provider,
router_correlation_id=litellm_call_id,
),
used_client_oauth_token=resolve_used_client_oauth_token(
metadata.get("used_client_oauth_token") if metadata is not None else None, custom_llm_provider
),
azure_spillover=azure_spillover(
response_headers=kwargs.get("response_headers")
if isinstance(kwargs.get("response_headers"), Mapping)
Expand Down
1 change: 1 addition & 0 deletions litellm/types/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3176,6 +3176,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
cold_storage_object_key: str | None # S3/GCS object key for cold storage retrieval
team_alias: str | None
team_id: str | None
used_client_oauth_token: ReadOnly[bool | None]


class AzureSpillover(TypedDict):
Expand Down
85 changes: 85 additions & 0 deletions tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,91 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m
assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info


@pytest.mark.asyncio
@pytest.mark.parametrize(
"used_client_oauth_token, custom_llm_provider, expected",
[(True, "anthropic", True), (True, "bedrock", False), (False, "anthropic", False)],
)
async def test_async_post_call_failure_hook_carries_used_client_oauth_token_from_litellm_metadata(
used_client_oauth_token: bool, custom_llm_provider: str, expected: bool
):
"""
/v1/messages and /v1/responses stamp the proxy's own fields into request_data["litellm_metadata"]
and leave request_data["metadata"] to the caller's native metadata, so a failed request on those
routes wrote a spend row whose used_client_oauth_token was null instead of the stamped value
"""
logger = _ProxyDBLogger()
request_data = {
"model": "claude-sonnet-5",
"custom_llm_provider": custom_llm_provider,
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {"user_id": "anthropic-native-metadata"},
"litellm_metadata": {"used_client_oauth_token": used_client_oauth_token},
"proxy_server_request": {"request_id": "test_request_id"},
}

with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("rate limited"),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
)

call_kwargs = mock_update_database.call_args[1]["kwargs"]
assert call_kwargs["litellm_params"]["metadata"]["user_id"] == "anthropic-native-metadata"
payload = get_logging_payload(
kwargs=call_kwargs, response_obj={}, start_time=datetime.now(), end_time=datetime.now()
)
assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected


@pytest.mark.asyncio
@pytest.mark.parametrize(
"metadata_buckets, expected",
[
({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"used_client_oauth_token": False}}, False),
({"metadata": {"used_client_oauth_token": True}, "litellm_metadata": {"user_id": "caller"}}, None),
({"metadata": {"used_client_oauth_token": "yes"}}, None),
],
)
async def test_async_post_call_failure_hook_never_lets_caller_metadata_set_used_client_oauth_token(
metadata_buckets: dict, expected: bool | None
):
"""
On /v1/messages and /v1/responses the request's own metadata field belongs to the caller, so a
used_client_oauth_token they put there must never outrank the proxy's stamp or stand in for a missing one
"""
logger = _ProxyDBLogger()
request_data = {
"model": "claude-sonnet-5",
"custom_llm_provider": "anthropic",
"messages": [{"role": "user", "content": "Hello"}],
"proxy_server_request": {"request_id": "test_request_id"},
**metadata_buckets,
}

with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("rate limited"),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
)

payload = get_logging_payload(
kwargs=mock_update_database.call_args[1]["kwargs"],
response_obj={},
start_time=datetime.now(),
end_time=datetime.now(),
)
assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected


@pytest.mark.asyncio
async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request():
"""LIT-5651: a request blocked by a guardrail never reaches the LLM, but the
Expand Down
Loading
Loading