A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CLI] Use streaming in CLI chat and completion commands (#23769)
Signed-off-by: simon-mo <simon.mo@hey.com>
Simon Mo committed
Sep 17, 2025 at 22:30 UTC
e111d5b0ae9359e2a829771105e739d36505fa69
1 file changed
+45
-26
vllm/entrypoints/cli/openai.py
+45
-26
@@ -45,6 +45,28 @@ def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]:
45
return model_name, openai_client
46
47
48
+def _print_chat_stream(stream) -> str:
49
+ output = ""
50
+ for chunk in stream:
51
+ delta = chunk.choices[0].delta
52
+ if delta.content:
53
+ output += delta.content
54
+ print(delta.content, end="", flush=True)
55
+ print()
56
+ return output
57
+
58
+
59
+def _print_completion_stream(stream) -> str:
60
+ output = ""
61
+ for chunk in stream:
62
+ text = chunk.choices[0].text
63
+ if text is not None:
64
+ output += text
65
+ print(text, end="", flush=True)
66
+ print()
67
+ return output
68
+
69
+
70
def chat(system_prompt: str | None, model_name: str, client: OpenAI) -> None:
71
conversation: list[ChatCompletionMessageParam] = []
72
if system_prompt is not None:
@@ -58,14 +80,11 @@ def chat(system_prompt: str | None, model_name: str, client: OpenAI) -> None:
80
break
81
conversation.append({"role": "user", "content": input_message})
82
61
- chat_completion = client.chat.completions.create(model=model_name,
62
- messages=conversation)
63
-
64
- response_message = chat_completion.choices[0].message
65
- output = response_message.content
66
-
67
- conversation.append(response_message) # type: ignore
68
- print(output)
83
+ stream = client.chat.completions.create(model=model_name,
84
+ messages=conversation,
85
+ stream=True)
86
+ output = _print_chat_stream(stream)
87
+ conversation.append({"role": "assistant", "content": output})
88
89
90
def _add_query_options(
@@ -108,9 +127,11 @@ class ChatCommand(CLISubcommand):
127
if args.quick:
128
conversation.append({"role": "user", "content": args.quick})
129
111
- chat_completion = client.chat.completions.create(
112
- model=model_name, messages=conversation)
113
- print(chat_completion.choices[0].message.content)
130
+ stream = client.chat.completions.create(model=model_name,
131
+ messages=conversation,
132
+ stream=True)
133
+ output = _print_chat_stream(stream)
134
+ conversation.append({"role": "assistant", "content": output})
135
return
136
137
print("Please enter a message for the chat model:")
@@ -121,14 +142,11 @@ class ChatCommand(CLISubcommand):
142
break
143
conversation.append({"role": "user", "content": input_message})
144
124
- chat_completion = client.chat.completions.create(
125
- model=model_name, messages=conversation)
126
-
127
- response_message = chat_completion.choices[0].message
128
- output = response_message.content
129
-
130
- conversation.append(response_message) # type: ignore
131
- print(output)
145
+ stream = client.chat.completions.create(model=model_name,
146
+ messages=conversation,
147
+ stream=True)
148
+ output = _print_chat_stream(stream)
149
+ conversation.append({"role": "assistant", "content": output})
150
151
@staticmethod
152
def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
@@ -168,9 +186,10 @@ class CompleteCommand(CLISubcommand):
186
model_name, client = _interactive_cli(args)
187
188
if args.quick:
171
- completion = client.completions.create(model=model_name,
172
- prompt=args.quick)
173
- print(completion.choices[0].text)
189
+ stream = client.completions.create(model=model_name,
190
+ prompt=args.quick,
191
+ stream=True)
192
+ _print_completion_stream(stream)
193
return
194
195
print("Please enter prompt to complete:")
@@ -179,10 +198,10 @@ class CompleteCommand(CLISubcommand):
198
input_prompt = input("> ")
199
except EOFError:
200
break
182
- completion = client.completions.create(model=model_name,
183
- prompt=input_prompt)
184
- output = completion.choices[0].text
185
- print(output)
201
+ stream = client.completions.create(model=model_name,
202
+ prompt=input_prompt,
203
+ stream=True)
204
+ _print_completion_stream(stream)
205
206
@staticmethod
207
def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: