mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
fix(proxy): carry used_client_oauth_token onto failure spend rows for litellm_metadata routes
This commit is contained in:
parent
23a4bc5b6c
commit
d688293c81
2 changed files with 48 additions and 7 deletions
|
|
@ -83,6 +83,10 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
|||
str(CallTypes.aretrieve_batch),
|
||||
)
|
||||
)
|
||||
_FAILURE_ROW_KEYS_LIFTED_FROM_LITELLM_METADATA: Final[tuple[str, ...]] = (
|
||||
"standard_logging_guardrail_information",
|
||||
"used_client_oauth_token",
|
||||
)
|
||||
|
||||
|
||||
def _proxy_spend_writer() -> DBSpendUpdateWriter:
|
||||
|
|
@ -195,13 +199,13 @@ class _ProxyDBLogger(CustomLogger):
|
|||
existing_metadata.update(_metadata)
|
||||
|
||||
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
|
||||
if (
|
||||
isinstance(litellm_metadata_bucket, dict)
|
||||
and "standard_logging_guardrail_information" not in existing_metadata
|
||||
):
|
||||
guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information")
|
||||
if guardrail_info is not None:
|
||||
existing_metadata["standard_logging_guardrail_information"] = guardrail_info
|
||||
existing_metadata.update(
|
||||
(key, litellm_metadata_bucket[key])
|
||||
for key in _FAILURE_ROW_KEYS_LIFTED_FROM_LITELLM_METADATA
|
||||
if isinstance(litellm_metadata_bucket, dict)
|
||||
and key not in existing_metadata
|
||||
and litellm_metadata_bucket.get(key) is not None
|
||||
)
|
||||
|
||||
if "litellm_params" not in request_data:
|
||||
request_data["litellm_params"] = {}
|
||||
|
|
|
|||
|
|
@ -160,6 +160,43 @@ 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", [True, False])
|
||||
async def test_async_post_call_failure_hook_carries_used_client_oauth_token_from_litellm_metadata(
|
||||
used_client_oauth_token: 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",
|
||||
"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 used_client_oauth_token
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue