From 96b5711a0ce33236893d4f93245f64b906fac81e Mon Sep 17 00:00:00 2001 From: Ruben Hakopian Date: Thu, 8 Feb 2024 17:25:28 -0800 Subject: [PATCH] google-vertexai[patch]: Fixed SafetySettings handling in streaming API in VertexAI (#17278) The streaming API doesn't separate safety_settings from the generation_config payload. As the result the following error is observed when using `stream` API. The functionality is correct with `invoke` API. The fix separates the `safety_settings` from params and sets it as argument to the `send_message` method. ``` ERROR: Unknown field for GenerationConfig: safety_settings Traceback (most recent call last): File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/langchain_core/language_models/chat_models.py", line 250, in stream raise e File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/langchain_core/language_models/chat_models.py", line 234, in stream for chunk in self._stream( File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/langchain_google_vertexai/chat_models.py", line 501, in _stream for response in responses: File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/vertexai/generative_models/_generative_models.py", line 921, in _send_message_streaming for chunk in stream: File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/vertexai/generative_models/_generative_models.py", line 514, in _generate_content_streaming request = self._prepare_request( ^^^^^^^^^^^^^^^^^^^^^^ File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/vertexai/generative_models/_generative_models.py", line 256, in _prepare_request gapic_generation_config = gapic_content_types.GenerationConfig( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/user/Library/Caches/pypoetry/virtualenvs/chatbot-worker-main-Ju-qIM-X-py3.12/lib/python3.12/site-packages/proto/message.py", line 576, in __init__ raise ValueError( ValueError: Unknown field for GenerationConfig: safety_settings ``` --------- Co-authored-by: Erick Friis --- .../langchain_google_vertexai/chat_models.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/partners/google-vertexai/langchain_google_vertexai/chat_models.py b/libs/partners/google-vertexai/langchain_google_vertexai/chat_models.py index 48709faebd..096f18006f 100644 --- a/libs/partners/google-vertexai/langchain_google_vertexai/chat_models.py +++ b/libs/partners/google-vertexai/langchain_google_vertexai/chat_models.py @@ -510,8 +510,13 @@ class ChatVertexAI(_VertexAICommon, BaseChatModel): # set param to `functions` until core tool/function calling implemented raw_tools = params.pop("functions") if "functions" in params else None tools = _format_tools_to_vertex_tool(raw_tools) if raw_tools else None + safety_settings = params.pop("safety_settings", None) responses = chat.send_message( - message, stream=True, generation_config=params, tools=tools + message, + stream=True, + generation_config=params, + safety_settings=safety_settings, + tools=tools, ) for response in responses: message = _parse_response_candidate(response.candidates[0])