From 9851fa7b1b3ca4203b43c0c94d362a2ffe7bec49 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 16 Aug 2024 11:35:43 -0700
Subject: [PATCH] return traces in bedrock guardrails when enabled

---
 litellm/utils.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/litellm/utils.py b/litellm/utils.py
index 414550be485..d915690fc6e 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -10226,6 +10226,14 @@ class CustomStreamWrapper:
                     return
             elif self.received_finish_reason is not None:
                 if self.sent_last_chunk is True:
+                    # Bedrock returns the guardrail trace in the last chunk - we want to return this here
+                    if (
+                        self.custom_llm_provider == "bedrock"
+                        and "trace" in model_response
+                    ):
+                        return model_response
+
+                    # Default - return StopIteration
                     raise StopIteration
                 # flush any remaining holding chunk
                 if len(self.holding_chunk) > 0: