community[patch]: add missing chunk parameter for _stream/_astream (#17807)

- Description: Add missing chunk parameter for _stream/_astream for some
chat models, make all chat models in a consistent behaviour.
- Issue: N/A
- Dependencies: N/A
This commit is contained in:
mackong
2024-02-22 07:32:28 +08:00
committed by GitHub
parent 1b0802babe
commit 31891092d8
15 changed files with 71 additions and 42 deletions

View File

@@ -154,9 +154,10 @@ class GigaChat(_BaseGigaChat, BaseChatModel):
for chunk in self._client.stream(payload):
if chunk.choices:
content = chunk.choices[0].delta.content
yield ChatGenerationChunk(message=AIMessageChunk(content=content))
cg_chunk = ChatGenerationChunk(message=AIMessageChunk(content=content))
yield cg_chunk
if run_manager:
run_manager.on_llm_new_token(content)
run_manager.on_llm_new_token(content, chunk=cg_chunk)
async def _astream(
self,
@@ -170,9 +171,10 @@ class GigaChat(_BaseGigaChat, BaseChatModel):
async for chunk in self._client.astream(payload):
if chunk.choices:
content = chunk.choices[0].delta.content
yield ChatGenerationChunk(message=AIMessageChunk(content=content))
cg_chunk = ChatGenerationChunk(message=AIMessageChunk(content=content))
yield cg_chunk
if run_manager:
await run_manager.on_llm_new_token(content)
await run_manager.on_llm_new_token(content, chunk=cg_chunk)
def get_num_tokens(self, text: str) -> int:
"""Count approximate number of tokens"""