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

@@ -328,9 +328,10 @@ class ChatDeepInfra(BaseChatModel):
for line in _parse_stream(response.iter_lines()):
chunk = _handle_sse_line(line)
if chunk:
yield ChatGenerationChunk(message=chunk, generation_info=None)
cg_chunk = ChatGenerationChunk(message=chunk, generation_info=None)
yield cg_chunk
if run_manager:
run_manager.on_llm_new_token(str(chunk.content))
run_manager.on_llm_new_token(str(chunk.content), chunk=cg_chunk)
async def _astream(
self,
@@ -350,9 +351,12 @@ class ChatDeepInfra(BaseChatModel):
async for line in _parse_stream_async(response.content):
chunk = _handle_sse_line(line)
if chunk:
yield ChatGenerationChunk(message=chunk, generation_info=None)
cg_chunk = ChatGenerationChunk(message=chunk, generation_info=None)
yield cg_chunk
if run_manager:
await run_manager.on_llm_new_token(str(chunk.content))
await run_manager.on_llm_new_token(
str(chunk.content), chunk=cg_chunk
)
async def _agenerate(
self,