From d0a78482018a7aa1c9645cda06745b284e7e370c Mon Sep 17 00:00:00 2001 From: Sergey Kozlov Date: Mon, 26 Feb 2024 19:30:51 +0600 Subject: [PATCH] Fix safety_attributes merging issue (#23) * Fix safety_attributes merging issue * Process safety_attributes only for non-gemini models --------- Co-authored-by: Sergey Kozlov --- libs/vertexai/langchain_google_vertexai/_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/libs/vertexai/langchain_google_vertexai/_utils.py b/libs/vertexai/langchain_google_vertexai/_utils.py index 5b8c59a3..641721af 100644 --- 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