Fix safety_attributes merging issue (#23)

* Fix safety_attributes merging issue * Process safety_attributes only for non-gemini models --------- Co-authored-by: Sergey Kozlov <[email protected]>
langchain-ai · Feb 26, 2024 · d0a7848 · d0a7848
1 parent 78d8845
commit d0a7848
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/libs/vertexai/langchain_google_vertexai/_utils.py b/libs/vertexai/langchain_google_vertexai/_utils.py
@@ -149,7 +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")
+
     return info