diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3f2e1d7301..0f97473312c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -457,8 +457,37 @@ def _pipeline_managed_guardrail_names( ) +def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None: + if not isinstance(modified_bucket_value, dict): + return + modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed + surviving_writes: Final = {key: value for key, value in modified_bucket.items() if key != "guardrails"} + existing_bucket: Final = data.get(bucket_key) + if isinstance(existing_bucket, dict): + cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed + else: + data[bucket_key] = surviving_writes + + +def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, object]) -> None: + """ + Copy metadata-bucket writes from a pipeline's working copy back onto the request. + + Post_call pipelines run step hooks against a copied request dict so the payload + already sent upstream stays untouched, but hooks record proxy-internal logging + state in the metadata buckets (``applied_guardrails`` for response headers, + ``standard_logging_guardrail_information`` for spend logs), and those writes + must reach the request dict the proxy keeps reading after the pipeline returns. + + The ``guardrails`` key is the executor's per-step activation flag for + ``should_run_guardrail``, not a hook write, so it stays in the working copy. + """ + for bucket_key in ("metadata", "litellm_metadata"): + _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) + + def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("stream") is not True: + if data.get("stream") is not True and data.get("background") is not True: return post_call_policies: Final = tuple( policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" @@ -470,9 +499,10 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None detail={ "error": { "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses yet: " - f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output " - "guardrails from pipeline steps to guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern streaming or background " + f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and " + "background=false, or move these policies' output guardrails from pipeline steps to " + "guardrails.add, which scans streamed output." ), "type": "guardrail_pipeline_error", "policies": list(post_call_policies), @@ -1667,11 +1697,16 @@ class ProxyLogging: Returns data dict if allowed, raises on block/modify_response. ``original_response`` is set on the post_call path, where the request payload (already sent upstream) must stay untouched; a replacement - response carried in ``modified_data`` is adopted by the caller. + response carried in ``modified_data`` is adopted by the caller, and + metadata-bucket writes (applied guardrails, guardrail logging info) + are merged back so headers and spend logs still see them. """ if result.terminal_action == "allow": - if result.modified_data is not None and original_response is None: - data.update(result.modified_data) + if result.modified_data is not None: + if original_response is None: + data.update(result.modified_data) + else: + _merge_pipeline_metadata_writes(data, result.modified_data) return data if result.terminal_action == "block": diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 4e1ccf71c5c..34a25d75bc0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( @@ -1139,18 +1140,132 @@ def test_handle_pipeline_result_modify_response_carries_original_response(): assert info.value.original_response is response -def test_handle_pipeline_result_allow_discards_modifications_on_post_call(): +def test_handle_pipeline_result_allow_on_post_call_keeps_metadata_writes_only(): data = {"a": 1, "metadata": {"guardrails": ["other"]}} result = MagicMock() result.terminal_action = "allow" - result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()} + result.modified_data = { + "a": 2, + "metadata": {"guardrails": ["other"], "applied_guardrails": ["gr-post"]}, + "response": object(), + } out = ProxyLogging._handle_pipeline_result( result=result, data=data, policy_name="p", original_response=litellm.ModelResponse() ) assert out is data - assert data == {"a": 1, "metadata": {"guardrails": ["other"]}} + assert data["a"] == 1 + assert "response" not in data + assert data["metadata"] == {"guardrails": ["other"], "applied_guardrails": ["gr-post"]} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_guardrail_metadata_writes_reach_request_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class HeaderWritingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"verdict": "pass"}, + request_data=data, + guardrail_status="success", + ) + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [HeaderWritingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + slg_entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_entries) == 1 + assert slg_entries[0]["guardrail_name"] == "gr-post" + + +@pytest.mark.asyncio +async def test_post_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + seen["count"] += 1 + return data + + pre_call_pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gr-pre", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-pre", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-pre"}, + }, + } + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" + ) + + assert seen["count"] == 1 @pytest.mark.asyncio @@ -1173,6 +1288,26 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( assert "stream=false" in info.value.detail["error"]["message"] +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(background=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="aresponses", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ["response-governance"] + assert "background=false" in info.value.detail["error"]["message"] + + def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) @@ -1183,6 +1318,12 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c ) is None ) + assert ( + _raise_for_streaming_post_call_pipelines( + {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + ) + is None + ) assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None assert ( _raise_for_streaming_post_call_pipelines( @@ -1191,3 +1332,4 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c is None ) assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None + assert _raise_for_streaming_post_call_pipelines({"background": True}) is None