From 96af8f31ae3654e0a889ad78a47c0265c7112b54 Mon Sep 17 00:00:00 2001
From: maang-h <55082429+maang-h@users.noreply.github.com>
Date: Tue, 2 Jul 2024 02:12:24 +0800
Subject: [PATCH] community[patch]: Invoke callback prior to yielding token
 (#23638)

- **Description:** Invoke callback prior to yielding token in stream and
astream methods for ChatZhipuAI.
- **Issue:** the issue #16913
---
 libs/community/langchain_community/chat_models/zhipuai.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/libs/community/langchain_community/chat_models/zhipuai.py b/libs/community/langchain_community/chat_models/zhipuai.py
index 9ec8ffa728c..f3573129dd7 100644
--- a/libs/community/langchain_community/chat_models/zhipuai.py
+++ b/libs/community/langchain_community/chat_models/zhipuai.py
@@ -496,9 +496,10 @@ class ChatZhipuAI(BaseChatModel):
                     chunk = ChatGenerationChunk(
                         message=chunk, generation_info=generation_info
                     )
-                    yield chunk
                     if run_manager:
                         run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+                    yield chunk
+
                     if finish_reason is not None:
                         break
 
@@ -582,8 +583,9 @@ class ChatZhipuAI(BaseChatModel):
                     chunk = ChatGenerationChunk(
                         message=chunk, generation_info=generation_info
                     )
-                    yield chunk
                     if run_manager:
                         await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+                    yield chunk
+
                     if finish_reason is not None:
                         break