Updating callback for invoking prior to yielding token

langchain-ai · Feb 6, 2024 · d562e2d · d562e2d
1 parent f027696
commit d562e2d
Showing 1 changed file with 3 additions and 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
@@ -573,9 +573,9 @@ def _stream(
                 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,
@@ -599,9 +599,10 @@ async def _astream(
                 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,