Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

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")

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(