From d8fd9a20ed40a5ff4505778be4137d3076fd05e9 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 18 Mar 2026 00:42:01 +0100 Subject: [PATCH] =?UTF-8?q?fix(proxy):=20address=20Greptile=20review=20?= =?UTF-8?q?=E2=80=94=20streaming=20request=5Fdata,=20OCR=20backward=20comp?= =?UTF-8?q?at,=20test=20coverage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Pass request_data to end-of-stream process_output_streaming_response call - Restore inputs.update() in OCR handler for third-party guardrail providers - Add streaming end-to-end test for guardrail logging passthrough --- .../ocr/guardrail_translation/handler.py | 9 +- .../unified_guardrail/unified_guardrail.py | 1 + .../test_openai_moderation_streaming.py | 103 ++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py index 795ce1d7cae..7d3797a1dbe 100644 --- a/litellm/llms/mistral/ocr/guardrail_translation/handler.py +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -134,12 +134,17 @@ class OCRHandler(BaseTranslation): request_data = {} # Add user metadata if available - if "litellm_metadata" not in request_data and user_api_key_dict is not None: + if user_api_key_dict is not None: user_metadata = self.transform_user_api_key_dict_to_metadata( user_api_key_dict ) if user_metadata: - request_data["litellm_metadata"] = user_metadata + # Preserve original behavior: inject metadata into inputs for + # third-party guardrail providers that read it from there + inputs.update(user_metadata) # type: ignore + # Also store in request_data for the logging pipeline + if "litellm_metadata" not in request_data: + request_data["litellm_metadata"] = user_metadata guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 80b6f987f66..a1623121da5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -459,6 +459,7 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, + request_data=request_data, ) except HTTPException as e: if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index c77a5d07b3b..2595a1df7f1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -170,3 +170,106 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_openai_moderation_streaming_end_of_stream_request_data_passthrough(): + """Test that streaming end-of-stream guardrail info flows through to the + real request_data (Bug 1 fix for streaming path).""" + from litellm.types.llms.openai import ( + OpenAIModerationResponse, + OpenAIModerationResult, + ) + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = OpenAIModerationResponse( + id="modr-stream-test", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=False, + categories={"hate": False, "violence": False}, + category_scores={"hate": 0.001, "violence": 0.002}, + category_applied_input_types={"hate": [], "violence": []}, + ) + ], + ) + + async def mock_stream(): + import litellm + + chunks_data = ["Hello", " world"] + for i, content in enumerate(chunks_data): + chunk = MagicMock(spec=ModelResponseStream) + chunk.model = "gpt-4" + choice = MagicMock() + choice.delta = MagicMock() + choice.delta.content = content + choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None + chunk.choices = [choice] + yield chunk + + import litellm + + mock_model_response = ModelResponse( + id="mock-stream-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", content="Hello world" + ), + finish_reason="stop", + ) + ], + ) + + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": { + "guardrails": ["test-openai-moderation"], + "guardrail_config": {"streaming_sampling_rate": 1}, + }, + } + + with patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + # Verify guardrail info reached the REAL request_data (not a throwaway) + guardrail_info_list = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + assert guardrail_info_list is not None, ( + "Guardrail info should be in request_data after streaming" + ) + info = guardrail_info_list[0] + assert info["guardrail_status"] == "success" + + # Full moderation response dict, NOT the simplified "allow" string + guardrail_resp = info["guardrail_response"] + assert isinstance(guardrail_resp, dict), ( + f"Expected full moderation response dict, got {type(guardrail_resp)}: {guardrail_resp}" + ) + assert "results" in guardrail_resp