From 3e5da911e4fd031b44ff40110abec2da3015e390 Mon Sep 17 00:00:00 2001 From: Leonid Kuligin Date: Tue, 20 Feb 2024 21:29:42 +0100 Subject: [PATCH] fixes after review --- .../langchain_google_vertexai/chat_models.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py index a33d0764..ab620b91 100644 --- a/libs/vertexai/langchain_google_vertexai/chat_models.py +++ b/libs/vertexai/langchain_google_vertexai/chat_models.py @@ -397,14 +397,14 @@ def _generate( ) generations = [ ChatGeneration( - message=_parse_response_candidate(c), + message=_parse_response_candidate(candidate), generation_info=get_generation_info( - c, + candidate, self._is_gemini_model, usage_metadata=response.to_dict().get("usage_metadata"), ), ) - for c in response.candidates + for candidate in response.candidates ] else: question = _get_question(messages) @@ -416,14 +416,14 @@ def _generate( response = chat.send_message(question.content, **msg_params) generations = [ ChatGeneration( - message=AIMessage(content=r.text), + message=AIMessage(content=candidate.text), generation_info=get_generation_info( - r, + candidate, self._is_gemini_model, usage_metadata=response.raw_prediction_response.metadata, ), ) - for r in response.candidates + for candidate in response.candidates ] return ChatResult(generations=generations)