fix: Avoid error when we have just the tool_calls in input (#17753)

* fix: Avoid error when we have just the tool_calls in input

* fix: Remove the tool call validation

* feat: Add unit test
This commit is contained in:
Lucas Sugi
2025-12-09 22:59:59 -08:00
committed by GitHub
parent b0a5a4b81d
commit c7fd8fabdb
2 changed files with 77 additions and 23 deletions
@@ -1287,36 +1287,38 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
filtered_messages = filter_result.payload_messages or mock_messages
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
)
if bedrock_response.get("action") == "BLOCKED":
raise Exception(
f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}"
# Bedrock will throw an error if there is no text to process
if filtered_messages:
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
)
# Apply any masking that was applied by the guardrail
if bedrock_response.get("action") == "BLOCKED":
raise Exception(
f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}"
)
output_list = bedrock_response.get("output")
if output_list:
# If the guardrail returned modified content, use that
for output_item in output_list:
text_content = output_item.get("text")
if text_content:
masked_text = str(text_content)
masked_texts.append(masked_text)
else:
outputs_list = bedrock_response.get("outputs")
if outputs_list:
# Fallback to outputs field if output is not available
for output_item in outputs_list:
# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")
if output_list:
# If the guardrail returned modified content, use that
for output_item in output_list:
text_content = output_item.get("text")
if text_content:
masked_text = str(text_content)
masked_texts.append(masked_text)
else:
outputs_list = bedrock_response.get("outputs")
if outputs_list:
# Fallback to outputs field if output is not available
for output_item in outputs_list:
text_content = output_item.get("text")
if text_content:
masked_text = str(text_content)
masked_texts.append(masked_text)
# If no output/outputs were provided, use the original texts
# This happens when the guardrail allows content without modification
@@ -1049,3 +1049,55 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch
), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}"
print(f"Parameter precedence test passed. URL: {prepped_request.url}")
@pytest.mark.asyncio
async def test_bedrock_apply_guardrail_with_only_tool_calls_response():
"""Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API"""
# Create a BedrockGuardrail instance
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
)
# Mock the make_bedrock_api_request method
with patch.object(
guardrail, "make_bedrock_api_request", new_callable=AsyncMock
) as mock_api_request:
# Test the apply_guardrail method with tool_calls in response
inputs = {
"texts": [],
"tool_calls": [
{
"id": "call_eFSCWFsyL7MclHYnzKrcQnMK",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location":"São Paulo"}',
},
}
],
}
guardrailed_inputs = await guardrail.apply_guardrail(
inputs=inputs,
request_data={},
input_type="response",
logging_obj=None,
)
# Verify the result - should succeed without errors
assert guardrailed_inputs is not None
assert "tool_calls" in guardrailed_inputs
assert len(guardrailed_inputs["tool_calls"]) == 1
assert (
guardrailed_inputs["tool_calls"][0]["id"]
== "call_eFSCWFsyL7MclHYnzKrcQnMK"
)
assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather"
assert (
guardrailed_inputs["tool_calls"][0]["function"]["arguments"]
== '{"location":"São Paulo"}'
)
# Verify that the Bedrock API was NOT called since there's no text to process
mock_api_request.assert_not_called()
print("✅ apply_guardrail with tool_calls test passed - no API call made")