fix(guardrails): skip INPUT validation when no user messages with experimental_use_latest_role_message_only#25355
Conversation
…erimental_use_latest_role_message_only Fixes BerriAI#23476
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes a bug where Key changes:
Confidence Score: 5/5Safe to merge — the fix is correct and all remaining findings are minor style issues that do not affect correctness. All findings are P2: redundant no-op assignments (cosmetic) and a missing streaming-path test (the logic is structurally identical to the tested success-hook path). The core bug fix is correct, existing tests pass, and no regressions are introduced. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py | Fix is correct: replaces the unsafe or fallback with explicit flag-aware logic in async_post_call_success_hook, async_post_call_streaming_iterator_hook, and apply_guardrail; adds a defense-in-depth role filter in _create_bedrock_input_content_request. Minor: three redundant = None no-op assignments. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py | Four new unit tests covering the role-guard filter, flag-disabled baseline, None-return from _prepare_guardrail_messages_for_role, and INPUT-skip in async_post_call_success_hook. All mocked — no real network calls. Streaming iterator hook path is untested. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[async_post_call_success_hook] --> B{should_validate_input?}
B -- No --> C[OUTPUT validation only]
B -- Yes --> D[_prepare_guardrail_messages_for_role]
D --> E{payload_messages is None?}
E -- No --> F[Run INPUT + OUTPUT in parallel]
E -- Yes --> G{experimental_use_latest_role_message_only?}
G -- No --> H[Fallback: use new_messages - Run INPUT + OUTPUT in parallel]
G -- Yes --> I[Skip INPUT entirely - Run OUTPUT only]
F --> J[Apply masking / blocked response]
H --> J
C --> J
I --> J
Reviews (1): Last reviewed commit: "fix(guardrails): skip INPUT validation w..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if input_messages is None: | ||
| if self.experimental_use_latest_role_message_only: | ||
| input_messages = None # no user messages → skip INPUT validation | ||
| else: | ||
| input_messages = new_messages |
There was a problem hiding this comment.
None reassignment (appears in 3 places)
Inside the if input_messages is None: guard, input_messages is already None, so input_messages = None on line 981 is a no-op and just adds noise. The same pattern appears at line 1138 in async_post_call_streaming_iterator_hook and line 1448 in apply_guardrail. Consider simplifying to make the intent clearer:
| if input_messages is None: | |
| if self.experimental_use_latest_role_message_only: | |
| input_messages = None # no user messages → skip INPUT validation | |
| else: | |
| input_messages = new_messages | |
| input_messages = input_filter.payload_messages | |
| if input_messages is None and not self.experimental_use_latest_role_message_only: | |
| input_messages = new_messages |
This removes the dead inner branch while keeping exactly the same runtime behavior.
| @pytest.mark.asyncio | ||
| async def test_post_call_success_hook_skips_input_when_no_user_messages_and_flag_enabled(): | ||
| """When experimental_use_latest_role_message_only is True and there are no | ||
| user messages, async_post_call_success_hook should skip INPUT validation | ||
| and only run OUTPUT validation.""" | ||
| guardrail = BedrockGuardrail( | ||
| guardrailIdentifier="test-id", | ||
| guardrailVersion="DRAFT", | ||
| experimental_use_latest_role_message_only=True, | ||
| ) | ||
|
|
||
| mock_user_api_key_dict = UserAPIKeyAuth() | ||
|
|
||
| mock_response = litellm.ModelResponse( | ||
| id="test-id", | ||
| choices=[ | ||
| litellm.Choices( | ||
| index=0, | ||
| message=litellm.Message(role="assistant", content="safe response"), | ||
| finish_reason="stop", | ||
| ) | ||
| ], | ||
| created=1234567890, | ||
| model="gpt-4o", | ||
| object="chat.completion", | ||
| ) | ||
|
|
||
| request_data = { | ||
| "model": "gpt-4o", | ||
| "messages": [ | ||
| {"role": "assistant", "content": "previous response"}, | ||
| {"role": "tool", "content": "tool result"}, | ||
| ], | ||
| } | ||
|
|
||
| call_sources = [] | ||
|
|
||
| async def mock_make_bedrock_api_request(source=None, **kwargs): | ||
| call_sources.append(source) | ||
| mock_bedrock = MagicMock() | ||
| mock_bedrock.get.return_value = None | ||
| return mock_bedrock | ||
|
|
||
| with patch.object( | ||
| guardrail, | ||
| "make_bedrock_api_request", | ||
| side_effect=mock_make_bedrock_api_request, | ||
| ): | ||
| await guardrail.async_post_call_success_hook( | ||
| data=request_data, | ||
| user_api_key_dict=mock_user_api_key_dict, | ||
| response=mock_response, | ||
| ) | ||
|
|
||
| # Only OUTPUT should have been validated, not INPUT | ||
| assert "OUTPUT" in call_sources | ||
| assert "INPUT" not in call_sources |
There was a problem hiding this comment.
async_post_call_streaming_iterator_hook received the same INPUT-skip logic (around line 1132) as async_post_call_success_hook, but there is no corresponding test verifying that it also skips INPUT when there are no user messages and the flag is enabled. Given that the two paths are structurally identical, adding a parallel test for the streaming path would be a low-effort safeguard against a future regression in that branch.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
1 participant
Fixes #23476
Relevant issues
Fixes #23476
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🐛 Bug Fix
Changes
When
experimental_use_latest_role_message_onlyis enabled and no user-role messages exist in the conversation, the post-call hooks (async_post_call_success_hook,async_post_call_streaming_iterator_hook) andapply_guardrailfall back to sending ALL messages — includingtoolandassistantrole content — to the Bedrock guardrail API. This happens because_prepare_guardrail_messages_for_rolereturnsNoneforpayload_messages, and theorfallback (payload_messages or all_messages) evaluates to the full unfiltered list.Fix: Replace the
orfallback in 3 locations with explicit logic that skips INPUT validation entirely when no user messages are found (runs OUTPUT validation only). Added a defense-in-depth role guard in_create_bedrock_input_content_requestthat filters out non-user messages when the flag is enabled.Tests: 4 new unit tests in
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py:_prepare_guardrail_messages_for_rolereturns None when no user messagesasync_post_call_success_hookskips INPUT and only runs OUTPUT when no user messages