Add Support for GuardRailConverseContentBlock (#557)

We introduce a new flag, guard_last_turn_only, which applies the
selective guardrail for only the last user message. For advanced use
cases, we also allow the dev to pass in raw (Bedrock) message blocks.

Issue Link: #540

---------

Co-authored-by: Michael Chin <[email protected]>

Author: BumaldaOverTheWater94

libs/aws/langchain_aws/chat_models/bedrock_converse.py

@@ -456,6 +456,15 @@ class Joke(BaseModel):
     request_metadata: Optional[Dict[str, str]] = None
     """Key-Value pairs that you can use to filter invocation logs."""
 
+    guard_last_turn_only: bool = False
+    """Boolean flag for applying the guardrail to only the last turn."""
+
+    raw_blocks: Optional[List[Dict[str, Any]]] = None
+    """Raw Bedrock message blocks that can be passed in.
+    LangChain will relay them unchanged, enabling any combination of content block types.
+    This is useful for custom guardrail wrapping
+    """
+
     model_config = ConfigDict(
         extra="forbid",
         populate_by_name=True,
@@ -634,11 +643,31 @@ def validate_environment(self) -> Self:
                 service_name="bedrock-runtime",
             )
 
+        if self.guard_last_turn_only and not self.guardrail_config:
+            raise ValueError(
+                "`guard_last_turn_only=True` but no `guardrail_config` supplied. "
+                "Provide a guardrail via `guardrail_config` or "
+                "disable `guard_last_turn_only`."
+            )
         return self
 
     def _get_base_model(self) -> str:
         return self.base_model_id if self.base_model_id else self.model_id
 
+    def _apply_guard_last_turn_only(self, messages: List[Dict[str, Any]]) -> None:
+        for msg in reversed(messages):
+            if msg.get("role") == "user":
+                new_content = []
+                for block in msg["content"]:
+                    if "text" in block:
+                        new_content.append(
+                            {"guardContent": {"text": {"text": block["text"]}}}
+                        )
+                    else:
+                        new_content.append(block)
+                msg["content"] = new_content
+                break
+
     def _generate(
         self,
         messages: List[BaseMessage],
@@ -647,7 +676,16 @@ def _generate(
         **kwargs: Any,
     ) -> ChatResult:
         """Top Level call"""
-        bedrock_messages, system = _messages_to_bedrock(messages)
+
+        if self.raw_blocks is not None:
+            logger.debug(f"Using raw blocks: {self.raw_blocks}")
+            bedrock_messages, system = self.raw_blocks, []
+        else:
+            bedrock_messages, system = _messages_to_bedrock(messages)
+            if self.guard_last_turn_only:
+                logger.debug("Applying selective guardrail to only the last turn")
+                self._apply_guard_last_turn_only(bedrock_messages)
+
         logger.debug(f"input message to bedrock: {bedrock_messages}")
         logger.debug(f"System message to bedrock: {system}")
         params = self._converse_params(
@@ -673,7 +711,15 @@ def _stream(
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
-        bedrock_messages, system = _messages_to_bedrock(messages)
+        if self.raw_blocks is not None:
+            logger.debug(f"Using raw blocks: {self.raw_blocks}")
+            bedrock_messages, system = self.raw_blocks, []
+        else:
+            bedrock_messages, system = _messages_to_bedrock(messages)
+            if self.guard_last_turn_only:
+                logger.debug("Applying selective guardrail to only the last turn")
+                self._apply_guard_last_turn_only(bedrock_messages)
+
         params = self._converse_params(
             stop=stop,
             **_snake_to_camel_keys(

libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py

@@ -1307,3 +1307,79 @@ def test_model_kwargs() -> None:
     )
     assert llm.additional_model_request_fields == {"temperature": 0.2}
     assert llm.temperature is None
+
+
+def _create_mock_llm_guard_last_turn_only() -> (
+    Tuple[ChatBedrockConverse, mock.MagicMock]
+):
+    """Utility to create an LLM with guard_last_turn_only=True and a mocked client."""
+    mocked_client = mock.MagicMock()
+    llm = ChatBedrockConverse(
+        client=mocked_client,
+        model="anthropic.claude-3-sonnet-20240229-v1:0",
+        region_name="us-west-2",
+        guard_last_turn_only=True,
+        guardrails={"guardrailId": "dummy-guardrail", "guardrailVersion": "1"},
+    )
+    return llm, mocked_client
+
+
+def test_guard_last_turn_only_no_guardrail_config() -> None:
+    """Test that an error is raised if guard_last_turn_only is True but no guardrail_config is provided."""
+    with pytest.raises(ValueError):
+        ChatBedrockConverse(
+            client=mock.MagicMock(),
+            model="anthropic.claude-3-sonnet-20240229-v1:0",
+            region_name="us-west-2",
+            guard_last_turn_only=True,
+        )
+
+
+def test_generate_guard_last_turn_only() -> None:
+    """Test that _generate() wraps ONLY the final user turn with guardContent."""
+    llm, mocked_client = _create_mock_llm_guard_last_turn_only()
+
+    mocked_client.converse.return_value = {
+        "output": {"message": {"content": [{"text": "ok"}]}},
+        "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
+    }
+
+    messages = [
+        HumanMessage(content="First user message"),
+        AIMessage(content="Assistant reply"),
+        HumanMessage(content="Second user message"),
+    ]
+
+    llm.invoke(messages)
+    _, kwargs = mocked_client.converse.call_args
+    bedrock_msgs = kwargs["messages"]
+
+    assert bedrock_msgs[0]["content"][0] == {"text": "First user message"}
+    # Last user turn is wrapped in guardContent
+    assert bedrock_msgs[-1]["content"][0] == {
+        "guardContent": {"text": {"text": "Second user message"}}
+    }
+
+
+def test_stream_guard_last_turn_only() -> None:
+    """Test that stream() applies guardContent to final user turn."""
+    llm, mocked_client = _create_mock_llm_guard_last_turn_only()
+
+    mocked_client.converse_stream.return_value = {
+        "stream": [{"messageStart": {"role": "assistant"}}]
+    }
+
+    messages = [
+        HumanMessage(content="Hello"),
+        AIMessage(content="Hi!"),
+        HumanMessage(content="How are you?"),
+    ]
+    list(llm.stream(messages))
+
+    _, kwargs = mocked_client.converse_stream.call_args
+    bedrock_msgs = kwargs["messages"]
+
+    assert bedrock_msgs[0]["content"][0] == {"text": "Hello"}
+    assert bedrock_msgs[-1]["content"][0] == {
+        "guardContent": {"text": {"text": "How are you?"}}
+    }