main
py 72 lines 2.39 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """An example shows how to generate chat completions from reasoning models
4 like DeepSeekR1.
5
6 To run this example, you need to start the vLLM server with the reasoning
7 parser:
8
9 ```bash
10 vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
11 --reasoning-parser deepseek_r1
12 ```
13
14 Unlike openai_chat_completion_with_reasoning.py, this example demonstrates the
15 streaming chat completions feature.
16
17 The streaming chat completions feature allows you to receive chat completions
18 in real-time as they are generated by the model. This is useful for scenarios
19 where you want to display chat completions to the user as they are generated
20 by the model.
21
22 Remember to check content and reasoning exist in `ChatCompletionChunk`,
23 content may not exist leading to errors if you try to access it.
24 """
25
26 from openai import OpenAI
27
28 # Modify OpenAI's API key and API base to use vLLM's API server.
29 openai_api_key = "EMPTY"
30 openai_api_base = "http://localhost:8000/v1"
31
32 messages = [{"role": "user", "content": "9.11 and 9.8, which is greater?"}]
33
34
35 def main():
36 client = OpenAI(
37 api_key=openai_api_key,
38 base_url=openai_api_base,
39 )
40
41 models = client.models.list()
42 model = models.data[0].id
43
44 # ruff: noqa: E501
45 # For granite: add: `extra_body={"chat_template_kwargs": {"thinking": True}}`
46 stream = client.chat.completions.create(model=model, messages=messages, stream=True)
47
48 print("client: Start streaming chat completions...")
49 printed_reasoning = False
50 printed_content = False
51
52 for chunk in stream:
53 # Safely extract reasoning and content from delta,
54 # defaulting to None if attributes don't exist or are empty strings
55 reasoning = getattr(chunk.choices[0].delta, "reasoning", None) or None
56 content = getattr(chunk.choices[0].delta, "content", None) or None
57
58 if reasoning is not None:
59 if not printed_reasoning:
60 printed_reasoning = True
61 print("reasoning:", end="", flush=True)
62 print(reasoning, end="", flush=True)
63 elif content is not None:
64 if not printed_content:
65 printed_content = True
66 print("\ncontent:", end="", flush=True)
67 # Extract and print the content
68 print(content, end="", flush=True)
69
70
71 if __name__ == "__main__":
72 main()