community[patch]: Invoke callback prior to yielding token (#24917)

**Description: Invoke callback prior to yielding token in stream method
for chat_models .**
**Issue**: https://github.com/langchain-ai/langchain/issues/16913
#16913
This commit is contained in:
alexqiao 2024-08-01 21:19:55 +08:00 committed by GitHub
parent b2dd9ffaaf
commit 8eb0bdead3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -134,9 +134,9 @@ class ChatFriendli(BaseChatModel, BaseFriendli):
for chunk in stream:
delta = chunk.choices[0].delta.content
if delta:
yield ChatGenerationChunk(message=AIMessageChunk(content=delta))
if run_manager:
run_manager.on_llm_new_token(delta)
yield ChatGenerationChunk(message=AIMessageChunk(content=delta))
async def _astream(
self,
@ -152,9 +152,9 @@ class ChatFriendli(BaseChatModel, BaseFriendli):
async for chunk in stream:
delta = chunk.choices[0].delta.content
if delta:
yield ChatGenerationChunk(message=AIMessageChunk(content=delta))
if run_manager:
await run_manager.on_llm_new_token(delta)
yield ChatGenerationChunk(message=AIMessageChunk(content=delta))
def _generate(
self,