| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import os |
| 6 | import signal |
| 7 | import sys |
| 8 | import time |
| 9 | from typing import TYPE_CHECKING |
| 10 | |
| 11 | from openai import OpenAI |
| 12 | from openai.types.chat import ChatCompletionMessageParam |
| 13 | |
| 14 | from vllm.entrypoints.cli.types import CLISubcommand |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 18 | else: |
| 19 | FlexibleArgumentParser = argparse.ArgumentParser |
| 20 | |
| 21 | |
| 22 | def _register_signal_handlers(): |
| 23 | def signal_handler(sig, frame): |
| 24 | sys.exit(0) |
| 25 | |
| 26 | signal.signal(signal.SIGINT, signal_handler) |
| 27 | signal.signal(signal.SIGTSTP, signal_handler) |
| 28 | |
| 29 | |
| 30 | def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]: |
| 31 | _register_signal_handlers() |
| 32 | |
| 33 | base_url = args.url |
| 34 | api_key = args.api_key or os.environ.get("OPENAI_API_KEY", "EMPTY") |
| 35 | openai_client = OpenAI(api_key=api_key, base_url=base_url) |
| 36 | |
| 37 | if args.model_name: |
| 38 | model_name = args.model_name |
| 39 | else: |
| 40 | available_models = openai_client.models.list() |
| 41 | model_name = available_models.data[0].id |
| 42 | |
| 43 | print(f"Using model: {model_name}") |
| 44 | |
| 45 | return model_name, openai_client |
| 46 | |
| 47 | |
| 48 | def _print_chat_stream(stream, stats: bool = False) -> str: |
| 49 | output: str = "" |
| 50 | in_reasoning: bool = False |
| 51 | start: float = time.perf_counter() |
| 52 | ttft: float | None = None |
| 53 | completion_tokens: int = 0 |
| 54 | for chunk in stream: |
| 55 | if chunk.usage is not None: |
| 56 | completion_tokens = chunk.usage.completion_tokens |
| 57 | if not chunk.choices: |
| 58 | continue |
| 59 | delta = chunk.choices[0].delta |
| 60 | reasoning = getattr(delta, "reasoning", None) or getattr( |
| 61 | delta, "reasoning_content", None |
| 62 | ) |
| 63 | if ttft is None and (reasoning or delta.content): |
| 64 | ttft = time.perf_counter() - start |
| 65 | if reasoning: |
| 66 | if not in_reasoning: |
| 67 | print("<think>", flush=True) |
| 68 | in_reasoning = True |
| 69 | print(reasoning, end="", flush=True) |
| 70 | if delta.content: |
| 71 | if in_reasoning: |
| 72 | print("\n</think>", flush=True) |
| 73 | in_reasoning = False |
| 74 | output += delta.content |
| 75 | print(delta.content, end="", flush=True) |
| 76 | if in_reasoning: |
| 77 | print("\n</think>", end="", flush=True) |
| 78 | print() |
| 79 | if stats: |
| 80 | _print_metrics(start, ttft, completion_tokens) |
| 81 | return output |
| 82 | |
| 83 | |
| 84 | def _print_metrics(start: float, ttft: float | None, completion_tokens: int) -> None: |
| 85 | total_time = time.perf_counter() - start |
| 86 | if ttft is None or total_time <= 0: |
| 87 | return |
| 88 | print(f"{'TTFT:':<5} {ttft * 1000:.2f} ms") |
| 89 | print( |
| 90 | f"{'TPS:':<5} {completion_tokens / total_time:.2f} tokens/s " |
| 91 | f"({completion_tokens} tokens in {total_time:.2f}s)" |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | def _print_completion_stream(stream, stats: bool = False) -> str: |
| 96 | output = "" |
| 97 | start = time.perf_counter() |
| 98 | ttft: float | None = None |
| 99 | completion_tokens = 0 |
| 100 | for chunk in stream: |
| 101 | if chunk.usage is not None: |
| 102 | completion_tokens = chunk.usage.completion_tokens |
| 103 | if not chunk.choices: |
| 104 | continue |
| 105 | text = chunk.choices[0].text |
| 106 | if text: |
| 107 | if ttft is None: |
| 108 | ttft = time.perf_counter() - start |
| 109 | output += text |
| 110 | print(text, end="", flush=True) |
| 111 | print() |
| 112 | if stats: |
| 113 | _print_metrics(start, ttft, completion_tokens) |
| 114 | return output |
| 115 | |
| 116 | |
| 117 | def chat(system_prompt: str | None, model_name: str, client: OpenAI) -> None: |
| 118 | conversation: list[ChatCompletionMessageParam] = [] |
| 119 | if system_prompt is not None: |
| 120 | conversation.append({"role": "system", "content": system_prompt}) |
| 121 | |
| 122 | print("Please enter a message for the chat model:") |
| 123 | while True: |
| 124 | try: |
| 125 | input_message = input("> ") |
| 126 | except EOFError: |
| 127 | break |
| 128 | conversation.append({"role": "user", "content": input_message}) |
| 129 | |
| 130 | stream = client.chat.completions.create( |
| 131 | model=model_name, messages=conversation, stream=True |
| 132 | ) |
| 133 | output = _print_chat_stream(stream) |
| 134 | conversation.append({"role": "assistant", "content": output}) |
| 135 | |
| 136 | |
| 137 | def _add_query_options(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 138 | parser.add_argument( |
| 139 | "--url", |
| 140 | type=str, |
| 141 | default="http://localhost:8000/v1", |
| 142 | help="url of the running OpenAI-Compatible RESTful API server", |
| 143 | ) |
| 144 | parser.add_argument( |
| 145 | "--model-name", |
| 146 | type=str, |
| 147 | default=None, |
| 148 | help=( |
| 149 | "The model name used in prompt completion, default to " |
| 150 | "the first model in list models API call." |
| 151 | ), |
| 152 | ) |
| 153 | parser.add_argument( |
| 154 | "--api-key", |
| 155 | type=str, |
| 156 | default=None, |
| 157 | help=( |
| 158 | "API key for OpenAI services. If provided, this api key " |
| 159 | "will overwrite the api key obtained through environment variables." |
| 160 | " It is important to note that this option only applies to the " |
| 161 | "OpenAI-compatible API endpoints and NOT other endpoints that may " |
| 162 | "be present in the server. See the security guide in the vLLM docs " |
| 163 | "for more details." |
| 164 | ), |
| 165 | ) |
| 166 | return parser |
| 167 | |
| 168 | |
| 169 | class ChatCommand(CLISubcommand): |
| 170 | """The `chat` subcommand for the vLLM CLI.""" |
| 171 | |
| 172 | name = "chat" |
| 173 | |
| 174 | @staticmethod |
| 175 | def cmd(args: argparse.Namespace) -> None: |
| 176 | model_name, client = _interactive_cli(args) |
| 177 | system_prompt = args.system_prompt |
| 178 | stats = args.stats |
| 179 | conversation: list[ChatCompletionMessageParam] = [] |
| 180 | |
| 181 | if system_prompt is not None: |
| 182 | conversation.append({"role": "system", "content": system_prompt}) |
| 183 | |
| 184 | create_kwargs = {"model": model_name, "stream": True} |
| 185 | if stats: |
| 186 | create_kwargs["stream_options"] = {"include_usage": True} |
| 187 | |
| 188 | if args.quick is not None: |
| 189 | conversation.append({"role": "user", "content": args.quick}) |
| 190 | |
| 191 | stream = client.chat.completions.create( |
| 192 | messages=conversation, **create_kwargs |
| 193 | ) |
| 194 | output = _print_chat_stream(stream, stats) |
| 195 | conversation.append({"role": "assistant", "content": output}) |
| 196 | return |
| 197 | |
| 198 | print("Please enter a message for the chat model:") |
| 199 | while True: |
| 200 | try: |
| 201 | input_message = input("> ") |
| 202 | except EOFError: |
| 203 | break |
| 204 | conversation.append({"role": "user", "content": input_message}) |
| 205 | |
| 206 | stream = client.chat.completions.create( |
| 207 | messages=conversation, **create_kwargs |
| 208 | ) |
| 209 | output = _print_chat_stream(stream, stats) |
| 210 | conversation.append({"role": "assistant", "content": output}) |
| 211 | |
| 212 | @staticmethod |
| 213 | def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 214 | """Add CLI arguments for the chat command.""" |
| 215 | _add_query_options(parser) |
| 216 | parser.add_argument( |
| 217 | "--system-prompt", |
| 218 | type=str, |
| 219 | default=None, |
| 220 | help=( |
| 221 | "The system prompt to be added to the chat template, " |
| 222 | "used for models that support system prompts." |
| 223 | ), |
| 224 | ) |
| 225 | parser.add_argument( |
| 226 | "-q", |
| 227 | "--quick", |
| 228 | type=str, |
| 229 | metavar="MESSAGE", |
| 230 | help=("Send a single prompt as MESSAGE and print the response, then exit."), |
| 231 | ) |
| 232 | parser.add_argument( |
| 233 | "--stats", |
| 234 | action="store_true", |
| 235 | help="Print TTFT and TPS statistics after each response.", |
| 236 | ) |
| 237 | return parser |
| 238 | |
| 239 | def subparser_init( |
| 240 | self, subparsers: argparse._SubParsersAction |
| 241 | ) -> FlexibleArgumentParser: |
| 242 | parser = subparsers.add_parser( |
| 243 | "chat", |
| 244 | help="Generate chat completions via the running API server.", |
| 245 | description="Generate chat completions via the running API server.", |
| 246 | usage="vllm chat [options]", |
| 247 | ) |
| 248 | return ChatCommand.add_cli_args(parser) |
| 249 | |
| 250 | |
| 251 | class CompleteCommand(CLISubcommand): |
| 252 | """The `complete` subcommand for the vLLM CLI.""" |
| 253 | |
| 254 | name = "complete" |
| 255 | |
| 256 | @staticmethod |
| 257 | def cmd(args: argparse.Namespace) -> None: |
| 258 | model_name, client = _interactive_cli(args) |
| 259 | stats = args.stats |
| 260 | |
| 261 | kwargs = { |
| 262 | "model": model_name, |
| 263 | "stream": True, |
| 264 | } |
| 265 | if args.max_tokens is not None: |
| 266 | kwargs["max_tokens"] = args.max_tokens |
| 267 | if stats: |
| 268 | kwargs["stream_options"] = {"include_usage": True} |
| 269 | |
| 270 | if args.quick is not None: |
| 271 | stream = client.completions.create(prompt=args.quick, **kwargs) |
| 272 | _print_completion_stream(stream, stats) |
| 273 | return |
| 274 | |
| 275 | print("Please enter prompt to complete:") |
| 276 | while True: |
| 277 | try: |
| 278 | input_prompt = input("> ") |
| 279 | except EOFError: |
| 280 | break |
| 281 | stream = client.completions.create(prompt=input_prompt, **kwargs) |
| 282 | _print_completion_stream(stream, stats) |
| 283 | |
| 284 | @staticmethod |
| 285 | def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 286 | """Add CLI arguments for the complete command.""" |
| 287 | _add_query_options(parser) |
| 288 | parser.add_argument( |
| 289 | "--max-tokens", |
| 290 | type=int, |
| 291 | help="Maximum number of tokens to generate per output sequence.", |
| 292 | ) |
| 293 | parser.add_argument( |
| 294 | "-q", |
| 295 | "--quick", |
| 296 | type=str, |
| 297 | metavar="PROMPT", |
| 298 | help="Send a single prompt and print the completion output, then exit.", |
| 299 | ) |
| 300 | parser.add_argument( |
| 301 | "--stats", |
| 302 | action="store_true", |
| 303 | help="Print TTFT and TPS statistics after each response.", |
| 304 | ) |
| 305 | return parser |
| 306 | |
| 307 | def subparser_init( |
| 308 | self, subparsers: argparse._SubParsersAction |
| 309 | ) -> FlexibleArgumentParser: |
| 310 | parser = subparsers.add_parser( |
| 311 | "complete", |
| 312 | help=( |
| 313 | "Generate text completions based on the given prompt " |
| 314 | "via the running API server." |
| 315 | ), |
| 316 | description=( |
| 317 | "Generate text completions based on the given prompt " |
| 318 | "via the running API server." |
| 319 | ), |
| 320 | usage="vllm complete [options]", |
| 321 | ) |
| 322 | return CompleteCommand.add_cli_args(parser) |
| 323 | |
| 324 | |
| 325 | def cmd_init() -> list[CLISubcommand]: |
| 326 | return [ChatCommand(), CompleteCommand()] |