From adcfecdb67a3477a5196773d9e9e0c6c411316be Mon Sep 17 00:00:00 2001 From: Subhrajyoty Roy Date: Fri, 27 Sep 2024 18:11:13 +0530 Subject: [PATCH] community[patch]: callback before yield for textgen (#26929) **Description:** Moves callback to before yield for `_stream` and `_astream` function for the textgen model in the community llm package **Issue:** #16913 --- libs/community/langchain_community/llms/textgen.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/libs/community/langchain_community/llms/textgen.py b/libs/community/langchain_community/llms/textgen.py index a9a100d4026..33a4f74f976 100644 --- a/libs/community/langchain_community/llms/textgen.py +++ b/libs/community/langchain_community/llms/textgen.py @@ -335,14 +335,13 @@ class TextGen(LLM): text=result["text"], # type: ignore[call-overload, index] generation_info=None, ) + if run_manager: + run_manager.on_llm_new_token(token=chunk.text) yield chunk elif result["event"] == "stream_end": # type: ignore[call-overload, index] websocket_client.close() return - if run_manager: - run_manager.on_llm_new_token(token=chunk.text) - async def _astream( self, prompt: str, @@ -408,10 +407,9 @@ class TextGen(LLM): text=result["text"], # type: ignore[call-overload, index] generation_info=None, ) + if run_manager: + await run_manager.on_llm_new_token(token=chunk.text) yield chunk elif result["event"] == "stream_end": # type: ignore[call-overload, index] websocket_client.close() return - - if run_manager: - await run_manager.on_llm_new_token(token=chunk.text)