scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(guardrails): transfer guardrail evaluation metadata to spend logs on success path

Browse source 
On successful requests, guardrail evaluations run and store results in
request_data["litellm_metadata"]["standard_logging_guardrail_information"],
but the spend-log serialization reads from request_data["metadata"], so
guardrail evaluations were never reported in Request Logs.

The failure path already had this transfer; this adds the same logic to
the success path so guardrail evaluation info appears in both success
and failure spend logs.

Fixes LIT-6314.
This commit is contained in:
Yucheng Zhu 2026-08-31 18:26:15 -07:00
parent 3fadcd7155
commit 698f608ad6
2 changed files with 85 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -269,6 +269,17 @@ class _ProxyDBLogger(CustomLogger):
served_model_id=sl_object.get("model_id") if sl_object is not None else None,
router=get_llm_router(),
)
# LIT-6314: post-call guardrail evaluations run after SLP is built, so
# guardrail_information is missing from SLP. Populate it from post-call evals
# before spend-log write so reports show accurate guardrail results.
if sl_object is not None:
guardrail_info_from_hooks: Final = (
kwargs.get("litellm_metadata", {}).get("standard_logging_guardrail_information")
if isinstance(kwargs.get("litellm_metadata"), dict)
else None
)
if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"):
sl_object["guardrail_information"] = guardrail_info_from_hooks
if response_cost is not None:
user_api_key: Final = metadata.get("user_api_key", None)

132
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -1,4 +1,3 @@
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
@ -70,9 +69,7 @@ async def test_async_post_call_failure_hook():
# Check that metadata was properly updated
assert "litellm_params" in call_args["kwargs"]
assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {
"request_id": "test_request_id"
}
assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"}
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
assert metadata["user_api_key"] == "test_api_key"
assert metadata["status"] == "failure"
@ -336,9 +333,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails():
)
assert mock_invalidate_budget_reservation_counters.await_count == 1
assert (
mock_invalidate_budget_reservation_counters.await_args.kwargs[
"budget_reservation"
]
mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"]
is user_api_key_dict.budget_reservation
)
assert user_api_key_dict.budget_reservation["finalized"] is True
@ -433,36 +428,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
"entries": [{"counter_key": "spend:key:test_api_key"}],
}
assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None
assert (
_get_budget_reservation_from_metadata(
metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}
)
is None
)
assert (
_get_budget_reservation_from_metadata(
metadata={
"user_api_key_auth": UserAPIKeyAuth(
budget_reservation=budget_reservation
)
}
metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)}
)
== budget_reservation
)
assert (
_get_budget_reservation_from_metadata(
metadata={
"user_api_key_auth": dict(
UserAPIKeyAuth(budget_reservation=budget_reservation)
)
}
metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))}
)
== budget_reservation
)
assert (
_get_budget_reservation_from_metadata(
metadata={"user_api_key_budget_reservation": budget_reservation}
)
_get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation})
is budget_reservation
)
@ -470,9 +450,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
@pytest.mark.asyncio
async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails():
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
side_effect=Exception("db unavailable")
)
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable"))
increment_spend_counters = AsyncMock()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
@ -508,9 +486,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u
async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails():
proxy_logging_obj = MagicMock()
db_exception = RuntimeError("db unavailable")
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
side_effect=db_exception
)
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception)
increment_spend_counters = AsyncMock()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
@ -554,12 +530,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re
budget_reservation=budget_reservation,
)
assert mock_log_exception.call_count == 2
mock_log_exception.assert_any_call(
"Failed to release budget reservation after database update failed"
)
mock_log_exception.assert_any_call(
"Failed to invalidate budget reservation counters after release failed"
)
mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed")
mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed")
increment_spend_counters.assert_not_awaited()
@ -1101,10 +1073,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj
# standard_logging_object should have been propagated from logging obj
assert call_kwargs.get("standard_logging_object") is not None
assert (
call_kwargs["standard_logging_object"]["trace_id"]
== "trace-id-from-logging-obj"
)
assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj"
# litellm_trace_id should also be propagated as a fallback
assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj"
@ -1691,9 +1660,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
"metadata": {},
"proxy_server_request": {"request_id": "rid"},
"response_cost": 3.5e-05,
"combined_usage_object": Usage(
prompt_tokens=30, completion_tokens=1, total_tokens=31
),
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31),
}
with patch(
@ -1772,15 +1739,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
assert mock_increment.call_args.kwargs["team_id"] == "team-123"
assert mock_increment.call_args.kwargs["org_id"] == "org-456"
update_kwargs = (
mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs
)
update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs
assert update_kwargs["user_id"] == "mcp-user@example.com"
assert update_kwargs["team_id"] == "team-123"
assert (
kwargs["litellm_params"]["metadata"]["user_api_key_user_id"]
== "mcp-user@example.com"
)
assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com"
@pytest.mark.parametrize(
@ -1828,9 +1790,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect
],
)
@pytest.mark.asyncio
async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
call_type, expect_spend_log
):
async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log):
"""Regression for LIT-3782: a pass-through request with auth=false reaches the
cost callback with no key/user/team/end-user. Before the fix the spend-log
write was skipped and the request never appeared in request/usage logs. It
@ -1876,9 +1836,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
end_time=datetime.now(),
)
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (
1 if expect_spend_log else 0
)
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0)
class _FakeDeploymentLookup:
@ -1989,3 +1947,61 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un
)
assert charged == ("premium", "tier0")
@pytest.mark.asyncio
async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object():
"""
LIT-6314 regression: post-call guardrails append evaluations to
request_data["litellm_metadata"]["standard_logging_guardrail_information"],
but the standard_logging_object (built pre-post-call) lacks this field.
The callback must populate standard_logging_object["guardrail_information"]
from post-call evals before spend-log write so reports show results.
"""
logger = _ProxyDBLogger()
guardrail_info = [
{
"guardrail_name": "bedrock-guard",
"guardrail_status": "success",
"guardrail_cost": 0.0001,
}
]
sl_object = {"response_cost": 0.01}
kwargs = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"litellm_metadata": {"standard_logging_guardrail_information": guardrail_info},
"litellm_params": {
"metadata": {"user_api_key": "test_key"},
},
"call_type": CallTypes.completion.value,
"response_cost": 0.01,
"standard_logging_object": sl_object,
}
with (
patch(
"litellm.proxy.proxy_server.increment_spend_counters",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.proxy_server.update_cache",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
) as mock_proxy_logging,
):
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
mock_proxy_logging.failed_tracking_alert = AsyncMock()
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == 1
assert sl_object.get("guardrail_information") == guardrail_info