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 @@ -283,7 +283,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 @@ -5726,6 +5729,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 @@ -5740,6 +5744,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 @@ -5789,6 +5796,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 @@ -6512,6 +6523,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 @@ -6523,6 +6535,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 @@ -6797,6 +6810,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 @@ -4068,6 +4068,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 @@ -2162,7 +2163,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 @@ -3454,13 +3457,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 @@ -3481,6 +3484,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 @@ -2429,6 +2429,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 @@ -2838,6 +2847,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 @@ -43,6 +43,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 @@ -144,6 +145,7 @@ def _get_router_metadata_for_spend_log(
"autorouter_savings",
"autorouter_savings_estimate",
"autorouter_baseline_observation",
"used_client_oauth_token",
)
)

Expand All @@ -168,6 +170,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 @@ -212,6 +215,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 @@ -227,6 +231,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 @@ -694,6 +699,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 @@ -3160,6 +3160,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
33 changes: 33 additions & 0 deletions tests/test_litellm/litellm_core_utils/test_litellm_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -4846,6 +4846,39 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload["litellm_call_id"] == call_id


@pytest.mark.parametrize(
"client_sent_oauth_token, custom_llm_provider, expected",
[(True, "anthropic", True), (True, "bedrock", False), (False, "anthropic", False), (None, "anthropic", None)],
)
def test_get_standard_logging_object_payload_resolves_used_client_oauth_token_against_the_selected_provider(
logging_obj, client_sent_oauth_token: bool | None, custom_llm_provider: str, expected: bool | None
):
"""The proxy stamps whether the client presented an Anthropic OAuth bearer before routing, but the
bearer only reaches an Anthropic deployment, so the logged flag must follow the provider that was called."""
from datetime import datetime

from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload

request_metadata = {} if client_sent_oauth_token is None else {"used_client_oauth_token": client_sent_oauth_token}
now = datetime.now()
payload = get_standard_logging_object_payload(
kwargs={
"model": "claude-sonnet-5",
"messages": [],
"custom_llm_provider": custom_llm_provider,
"litellm_params": {"metadata": request_metadata},
},
init_response_obj={},
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="success",
)

assert payload is not None
assert payload["metadata"]["used_client_oauth_token"] is expected


def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj):
"""Access groups stamped at auth time reach the logging payload, so integrations see what a request billed."""
from datetime import datetime
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