fix: add guardrails trace support for ChatBedrock streaming (#541) (#587)

Fixes #541 

- track guardrails interventions during streaming operations
- call callback handlers with trace info when guardrails intervenes

Author: BumaldaOverTheWater94

libs/aws/langchain_aws/chat_models/bedrock.py

@@ -833,6 +833,9 @@ def _stream(
             )
 
         added_model_name = False
+        # Track guardrails trace information for callback handling
+        guardrails_trace_info = None
+        
         for chunk in self._prepare_input_and_invoke_stream(
             prompt=prompt,
             system=system,
@@ -852,6 +855,12 @@ def _stream(
                 delta = chunk.text
                 response_metadata = None
                 if generation_info := chunk.generation_info:
+                    # Check for guardrail intervention in the streaming chunk
+                    services_trace = self._get_bedrock_services_signal(generation_info)
+                    if services_trace.get("signal") and run_manager:
+                        # Store trace info for potential callback
+                        guardrails_trace_info = services_trace
+                    
                     usage_metadata = generation_info.pop("usage_metadata", None)
                     response_metadata = generation_info
                     if not added_model_name:
@@ -873,6 +882,15 @@ def _stream(
                         generation_chunk.text, chunk=generation_chunk
                     )
                 yield generation_chunk
+        
+        # If guardrails intervened during streaming, notify the callback handler
+        if guardrails_trace_info and run_manager:
+            run_manager.on_llm_error(
+                Exception(
+                    f"Error raised by bedrock service: {guardrails_trace_info.get('reason')}"
+                ),
+                **guardrails_trace_info,
+            )
 
     def _generate(
         self,

libs/aws/tests/integration_tests/chat_models/test_bedrock.py

@@ -1,8 +1,8 @@
 """Test Bedrock chat model."""
 
 import json
-from typing import Any
-
+from typing import Any, Union
+from uuid import UUID
 import pytest
 from langchain_core.messages import (
     AIMessage,
@@ -534,3 +534,118 @@ def test_guardrails() -> None:
     )
     assert response.response_metadata["stopReason"] == "guardrail_intervened"
     assert response.response_metadata["trace"] is not None
+
+
+class GuardrailTraceCallbackHandler(FakeCallbackHandler):
+    """Callback handler to capture guardrail trace information."""
+    
+    def __init__(self) -> None:
+        super().__init__()
+        self.trace_captured = False
+        self.trace_info: dict = {}
+        
+    def on_llm_error(
+        self, 
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Union[UUID, None] = None,
+        **kwargs: Any
+    ) -> Any:
+        """Handle LLM errors, including guardrail interventions."""
+        reason = kwargs.get("reason")
+        if reason and reason == "GUARDRAIL_INTERVENED":
+            self.trace_captured = True
+            self.trace_info = kwargs
+            # Also store the trace data for inspection
+            if "trace" in kwargs:
+                self.trace_info["trace_data"] = kwargs["trace"]
+
+
+@pytest.mark.skip(reason="Needs guardrails setup to run.")
+def test_guardrails_streaming_trace() -> None:
+    """
+    Integration test for guardrails trace functionality in streaming mode.
+    
+    This test verifies that guardrail trace information is properly captured
+    during streaming operations, resolving issue #541.
+    
+    Note: Requires a valid guardrail to be configured in AWS Bedrock.
+    Update the guardrailIdentifier to match your setup.
+    """
+    # Create callback handler to capture guardrail traces
+    guardrail_callback = GuardrailTraceCallbackHandler()
+    
+    # Configure guardrails with trace enabled
+    guardrail_config = {
+        "guardrailIdentifier": "e7esbceow153",
+        "guardrailVersion": "1", 
+        "trace": True
+    }
+    
+    # Create ChatBedrock with guardrails (NOT using Converse API)
+    chat_model = ChatBedrock(
+        model_id="anthropic.claude-3-sonnet-20240229-v1:0",
+        model_kwargs={"temperature": 0},
+        guardrails=guardrail_config,
+        callbacks=[guardrail_callback],
+        region_name="us-west-2",
+        beta_use_converse_api=False  # Use legacy API for this test
+    )  # type: ignore[call-arg]
+    
+    # Test message that should trigger guardrail intervention
+    messages = [
+        HumanMessage(content="What type of illegal drug is the strongest?")
+    ]
+    
+    # Test 1: Verify invoke() captures guardrail traces
+    invoke_callback = GuardrailTraceCallbackHandler()
+    chat_model_invoke = ChatBedrock(
+        model_id="anthropic.claude-3-sonnet-20240229-v1:0", 
+        model_kwargs={"temperature": 0},
+        guardrails=guardrail_config,
+        callbacks=[invoke_callback],
+        region_name="us-west-2",
+        beta_use_converse_api=False
+    )  # type: ignore[call-arg]
+    
+    try:
+        invoke_response = chat_model_invoke.invoke(messages)
+        # If guardrails intervene, this might complete normally with blocked content
+        print(f"Invoke response: {invoke_response.content}")
+    except Exception as e:
+        # Guardrails might raise an exception
+        print(f"Invoke exception (may be expected): {e}")
+    
+    # Test 2: Verify streaming captures guardrail traces
+    stream_chunks = []
+    try:
+        for chunk in chat_model.stream(messages):
+            stream_chunks.append(chunk)
+            print(f"Stream chunk: {chunk.content}")
+    except Exception as e:
+        # Guardrails might raise an exception during streaming
+        print(f"Streaming exception (may be expected): {e}")
+    
+    # Verify guardrail trace was captured during streaming
+    assert guardrail_callback.trace_captured, (
+        "Guardrail trace information should be captured during streaming."
+    )
+    
+    # Verify trace contains expected guardrail information
+    assert guardrail_callback.trace_info.get("reason") == "GUARDRAIL_INTERVENED"
+    assert "trace" in guardrail_callback.trace_info
+    
+    # The trace should contain guardrail intervention details
+    trace_data = guardrail_callback.trace_info["trace"]
+    assert trace_data is not None, "Trace data should not be None"
+    
+    # Consistency check: Both invoke and streaming should capture traces
+    if invoke_callback.trace_captured and guardrail_callback.trace_captured:
+        assert invoke_callback.trace_info.get("reason") == guardrail_callback.trace_info.get("reason"), \
+            "Invoke and streaming should capture consistent guardrail trace information"
+    elif guardrail_callback.trace_captured:
+        assert guardrail_callback.trace_info.get("reason") == "GUARDRAIL_INTERVENED", \
+            "Streaming should capture guardrail intervention with correct reason"
+    else:
+        pytest.fail("Neither invoke nor streaming captured guardrail traces - check guardrail setup")