community[patch]: callback before yield for titan takeoff

langchain-ai · Sep 27, 2024 · 893eb61 · 893eb61
1 parent c6350d6
commit 893eb61
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/libs/community/langchain_community/llms/titan_takeoff.py b/libs/community/langchain_community/llms/titan_takeoff.py
@@ -252,13 +252,13 @@ def _stream(
                 if buffer:  # Ensure that there's content to process.
                     chunk = GenerationChunk(text=buffer)
                     buffer = ""  # Reset buffer for the next set of data.
-                    yield chunk
                     if run_manager:
                         run_manager.on_llm_new_token(token=chunk.text)
+                    yield chunk
 
         # Yield any remaining content in the buffer.
         if buffer:
             chunk = GenerationChunk(text=buffer.replace("</s>", ""))
-            yield chunk
             if run_manager:
                 run_manager.on_llm_new_token(token=chunk.text)
+            yield chunk