diff --git a/libs/vertexai/langchain_google_vertexai/_utils.py b/libs/vertexai/langchain_google_vertexai/_utils.py
index 794c74d1..641721af 100644
--- a/libs/vertexai/langchain_google_vertexai/_utils.py
+++ b/libs/vertexai/langchain_google_vertexai/_utils.py
@@ -149,15 +149,16 @@ def get_generation_info(
             info_usage_metadata["prompt_token_count"] = input_usage.get("totalTokens")
             info["usage_metadata"] = {k: v for k, v in info_usage_metadata.items() if v}
 
+            # NOTE:
+            # "safety_attributes" can contain different values for the same keys
+            # for each generation. Put it in a list, so it can be merged later by
+            # merge_dicts().
+            #
+            safety_attributes = info.get("safety_attributes") or {}
+            info["safety_attributes"] = [safety_attributes]
+
     if stream:
         # Remove non-streamable types, like bools.
         info.pop("is_blocked")
 
-    # NOTE:
-    # "safety_attributes" can contain different values for the same keys
-    # for each generation. Put it in a list, so it can be merged later by merge_dicts().
-    #
-    safety_attributes = info.get("safety_attributes") or {}
-    info["safety_attributes"] = [safety_attributes]
-
     return info