From a2167614b7683a373b7acd53c3b58b2d8ce903d2 Mon Sep 17 00:00:00 2001 From: Sparsh Jain <31439850+dudesparsh@users.noreply.github.com> Date: Fri, 9 Feb 2024 02:43:46 +0530 Subject: [PATCH] google-genai[patch]: Invoke callback prior to yielding token (#17092) - **Description:** Invoke callback prior to yielding token in stream and astream methods for Google-genai, - **Issue:** the issue # 16913, - **Twitter handle:** Sparsh10649446 --------- Co-authored-by: Eugene Yurtsev --- .../google-genai/langchain_google_genai/chat_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/partners/google-genai/langchain_google_genai/chat_models.py b/libs/partners/google-genai/langchain_google_genai/chat_models.py index 2f0d35428d5..b69ff5824f0 100644 --- a/libs/partners/google-genai/langchain_google_genai/chat_models.py +++ b/libs/partners/google-genai/langchain_google_genai/chat_models.py @@ -624,9 +624,9 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel): generation_t=ChatGenerationChunk, ) gen = cast(ChatGenerationChunk, _chat_result.generations[0]) - yield gen if run_manager: run_manager.on_llm_new_token(gen.text) + yield gen async def _astream( self, @@ -654,9 +654,9 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel): generation_t=ChatGenerationChunk, ) gen = cast(ChatGenerationChunk, _chat_result.generations[0]) - yield gen if run_manager: await run_manager.on_llm_new_token(gen.text) + yield gen def _prepare_chat( self,