Skip to content

Commit

Permalink
fix vllm llamaindex stream bug (#679)
Browse files Browse the repository at this point in the history
Signed-off-by: Xinyao Wang <[email protected]>
  • Loading branch information
XinyaoWa authored Sep 12, 2024
1 parent 6ce6551 commit ca94c60
Showing 1 changed file with 3 additions and 5 deletions.
8 changes: 3 additions & 5 deletions comps/llms/text-generation/vllm/llama_index/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,13 +57,11 @@ def llm_generate(input: LLMParamsDoc):
if input.streaming:

def stream_generator():
chat_response = ""
for text in llm.stream_complete(input.query):
chat_response += text
chunk_repr = repr(text.encode("utf-8"))
yield f"data: {chunk_repr}\n\n"
output = text.text
yield f"data: {output}\n\n"
if logflag:
logger.info(f"[llm - chat_stream] stream response: {chat_response}")
logger.info(f"[llm - chat_stream] stream response: {output}")
yield "data: [DONE]\n\n"

return StreamingResponse(stream_generator(), media_type="text/event-stream")
Expand Down

0 comments on commit ca94c60

Please sign in to comment.