community/chat_models/bedrock.py - added yield after callback

langchain-ai · Dec 25, 2024 · 3d5514a · 3d5514a
1 parent 2a7469e
commit 3d5514a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/libs/community/langchain_community/chat_models/bedrock.py b/libs/community/langchain_community/chat_models/bedrock.py
@@ -264,7 +264,10 @@ def _stream(
             **kwargs,
         ):
             delta = chunk.text
-            yield ChatGenerationChunk(message=AIMessageChunk(content=delta))
+            cg_chunk =  ChatGenerationChunk(message=AIMessageChunk(content=delta))
+            if run_manager:
+                run_manager.on_llm_new_token(delta, chunk = cg_chunk )
+            yield cg_chunk
 
     def _generate(
         self,