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 15, 09:45 UTC
[Doc] Add stream flag for chat completion example (#18524)
Signed-off-by: calvin chen <120380290@qq.com>
Calvin Chen committed
May 22, 2025 at 22:07 UTC
3f505233fdd25e4026bbc1824eaa7a19aa782d80
1 file changed
+21
-3
examples/online_serving/openai_chat_completion_client.py
+21
-3
@@ -3,6 +3,9 @@
3
NOTE: start a supported chat completion model server with `vllm serve`, e.g.
4
vllm serve meta-llama/Llama-2-7b-chat-hf
5
"""
6
+
7
+import argparse
8
+
9
from openai import OpenAI
10
11
# Modify OpenAI's API key and API base to use vLLM's API server.
@@ -24,7 +27,15 @@ messages = [{
27
}]
28
29
27
-def main():
30
+def parse_args():
31
+ parser = argparse.ArgumentParser(description="Client for vLLM API server")
32
+ parser.add_argument("--stream",
33
+ action="store_true",
34
+ help="Enable streaming response")
35
+ return parser.parse_args()
36
+
37
+
38
+def main(args):
39
client = OpenAI(
40
# defaults to os.environ.get("OPENAI_API_KEY")
41
api_key=openai_api_key,
@@ -34,16 +45,23 @@ def main():
45
models = client.models.list()
46
model = models.data[0].id
47
48
+ # Chat Completion API
49
chat_completion = client.chat.completions.create(
50
messages=messages,
51
model=model,
52
+ stream=args.stream,
53
)
54
55
print("-" * 50)
56
print("Chat completion results:")
44
- print(chat_completion)
57
+ if args.stream:
58
+ for c in chat_completion:
59
+ print(c)
60
+ else:
61
+ print(chat_completion)
62
print("-" * 50)
63
64
65
if __name__ == "__main__":
49
- main()
66
+ args = parse_args()
67
+ main(args)