github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-22 06:39:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: callback before yield for textgen (#26929)

Browse Source 
**Description:** Moves callback to before yield for `_stream` and
`_astream` function for the textgen model in the community llm package
**Issue:** #16913
This commit is contained in:
Subhrajyoty Roy 2024-09-27 18:11:13 +05:30 committed by GitHub
parent 5f2cc4ecb2
commit adcfecdb67
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
libs/community/langchain_community/llms/textgen.py
View File

@ -335,14 +335,13 @@ class TextGen(LLM):
text=result["text"], # type: ignore[call-overload, index] text=result["text"], # type: ignore[call-overload, index]
generation_info=None, generation_info=None,
) )
if run_manager:
run_manager.on_llm_new_token(token=chunk.text)
yield chunk yield chunk
elif result["event"] == "stream_end": # type: ignore[call-overload, index] elif result["event"] == "stream_end": # type: ignore[call-overload, index]
websocket_client.close() websocket_client.close()
return return
if run_manager:
run_manager.on_llm_new_token(token=chunk.text)
async def _astream( async def _astream(
self, self,
prompt: str, prompt: str,
@ -408,10 +407,9 @@ class TextGen(LLM):
text=result["text"], # type: ignore[call-overload, index] text=result["text"], # type: ignore[call-overload, index]
generation_info=None, generation_info=None,
) )
if run_manager:
await run_manager.on_llm_new_token(token=chunk.text)
yield chunk yield chunk
elif result["event"] == "stream_end": # type: ignore[call-overload, index] elif result["event"] == "stream_end": # type: ignore[call-overload, index]
websocket_client.close() websocket_client.close()
return return
if run_manager:
await run_manager.on_llm_new_token(token=chunk.text)