diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 76049974610..a527c8f2a2d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -215,7 +215,12 @@ _FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( ("guardrail_config", (dict,)), ("proxy_server_request", (dict,)), ) -_POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags") +_POST_CALL_SELECTION_FIELDS: Final = ( + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "spend_logs_metadata", + "tags", +) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) _GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset( @@ -1152,8 +1157,9 @@ def _holds_resolved_pipelines(pipelines: object) -> bool: def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: - """The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail - mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" + """The policy pipelines resolved for a request that inherits auth guardrails, the tags a tag-scoped guardrail mode + reads and the spend_logs_metadata a blocked response's row keeps, from the logging metadata that request setup moved + the body's ``metadata`` into""" if not inherited: return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} @@ -1282,10 +1288,19 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict[str, ``standard_logging_guardrail_information`` there and not on the dict the failure handler forwards to ``post_call_failure_hook``. Without this the otel guardrail span is emitted on allow but missing on block. Carry the - entries over so the failure path matches the unified path. + entries over so the failure path matches the unified path. The failure row + reads its ``metadata`` here too, so it keeps the ``spend_logs_metadata`` a + built-in route's post-call pass was given, as its success row would. """ if guardrail_data is None: return + spend_logs_metadata: Final = (as_str_mapping(guardrail_data.get("metadata")) or {}).get("spend_logs_metadata") + if spend_logs_metadata is not None: + request_fields: Final = cast("Mapping[str, object]", request_data) # cast-ok: the JSON body plus logging kwargs + request_data["metadata"] = { # rebind-ok: the failure path hands this dict to post_call_failure_hook + **(as_str_mapping(request_fields.get("metadata")) or {}), + "spend_logs_metadata": spend_logs_metadata, + } source_key: Final = get_metadata_variable_name_from_kwargs(guardrail_data) source_metadata: Final = as_str_mapping(guardrail_data.get(source_key)) or {} entries: Final = source_metadata.get("standard_logging_guardrail_information") diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 86792e1ff28..80d64543449 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -791,6 +791,71 @@ def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( assert int(str(rows[0]["total_tokens"])) > 0, rows +def _spend_logs_metadata_of_the_failure_row(gateway: Gateway, key: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (key_hash(key),)), + bool, + ) + assert [row["status"] for row in rows] == ["failure"], rows + return _JSON.validate_python(rows[0]["metadata"]).get("spend_logs_metadata") + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_a_blocked_response_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"spend_logs_metadata": {"team_field": "team"}}) + key: Final = scenario.key( + team_id=team, + metadata={"guardrails": [BLOCKING_POST_CALL], "spend_logs_metadata": {"key_field": "key"}}, + ) + _drain(rig) + body: Final = { + **_route_body(route, "What is the capital of France?"), + "metadata": {"spend_logs_metadata": {"body_field": "body"}}, + } + + response: Final = _send(gateway, route, body, key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == { + "body_field": "body", + "key_field": "key", + "team_field": "team", + } + + +def test_a_response_a_pipeline_replaces_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"spend_logs_metadata": {"key_field": "key"}}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_MODIFIES), key) + + assert response.status_code == 200, response.text + assert b"content_filter" in response.content, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == {"key_field": "key"} + + +def test_a_custom_endpoints_blocked_response_row_gets_no_spend_logs_metadata(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), + metadata={ + "allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE], + "spend_logs_metadata": {"key_field": "key"}, + }, + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) is None + + def _row_fields(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: return tuple( eventually(