@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 194 lines 6.61 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Examples of batched chat completions via the vLLM OpenAI-compatible API.
4
5 The /v1/chat/completions/batch endpoint accepts ``messages`` as a list of
6 conversations. Each conversation is processed independently and the response
7 contains one choice per conversation, indexed 0, 1, ..., N-1.
8
9 Start a server first, e.g.:
10 vllm serve Qwen/Qwen2.5-1.5B-Instruct --port 8000
11
12 Current limitations compared to /v1/chat/completions:
13 - Streaming is not supported.
14 - Tool use is not supported.
15 - Beam search is not supported.
16 """
17
18 import json
19 import os
20
21 import httpx
22
23 BASE_URL = os.environ.get("VLLM_BASE_URL", "http://localhost:8000")
24 MODEL = os.environ.get("VLLM_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
25 BATCH_URL = f"{BASE_URL}/v1/chat/completions/batch"
26
27
28 def post_batch(payload: dict) -> dict:
29 response = httpx.post(BATCH_URL, json=payload, timeout=60)
30 response.raise_for_status()
31 return response.json()
32
33
34 def main() -> None:
35 print("=== Example 1a: single conversation (standard endpoint) ===")
36 response = httpx.post(
37 f"{BASE_URL}/v1/chat/completions",
38 json={
39 "model": MODEL,
40 "messages": [{"role": "user", "content": "What is the capital of Japan?"}],
41 },
42 timeout=60,
43 )
44 response.raise_for_status()
45 data = response.json()
46 for choice in data["choices"]:
47 print(f" [{choice['index']}] {choice['message']['content']}")
48
49 print("\n=== Example 1b: batched plain text (2 conversations) ===")
50 data = post_batch(
51 {
52 "model": MODEL,
53 "messages": [
54 [{"role": "user", "content": "What is the capital of France?"}],
55 [{"role": "user", "content": "What is the capital of Japan?"}],
56 ],
57 }
58 )
59 for choice in data["choices"]:
60 print(f" [{choice['index']}] {choice['message']['content']}")
61
62 print("\n=== Example 2: batch with regex constraint (yes|no) ===")
63 data = post_batch(
64 {
65 "model": MODEL,
66 "messages": [
67 [{"role": "user", "content": "Is the sky blue? Answer yes or no."}],
68 [{"role": "user", "content": "Is fire cold? Answer yes or no."}],
69 ],
70 "structured_outputs": {"regex": "(yes|no)"},
71 }
72 )
73 for choice in data["choices"]:
74 print(f" [{choice['index']}] {choice['message']['content']}")
75
76 print("\n=== Example 3: batch with json_schema ===")
77 person_schema = {
78 "type": "object",
79 "properties": {
80 "name": {"type": "string", "description": "Full name of the person"},
81 "age": {"type": "integer", "description": "Age in years"},
82 },
83 "required": ["name", "age"],
84 }
85 data = post_batch(
86 {
87 "model": MODEL,
88 "messages": [
89 [
90 {
91 "role": "user",
92 "content": "Describe the person: name Alice, age 30.",
93 }
94 ],
95 [{"role": "user", "content": "Describe the person: name Bob, age 25."}],
96 ],
97 "response_format": {
98 "type": "json_schema",
99 "json_schema": {
100 "name": "person",
101 "strict": True,
102 "schema": person_schema,
103 },
104 },
105 }
106 )
107 for choice in data["choices"]:
108 person = json.loads(choice["message"]["content"])
109 print(f" [{choice['index']}] {person}")
110
111 print("\n=== Example 4: batch book summaries ===")
112 book_schema = {
113 "type": "object",
114 "properties": {
115 "author": {
116 "type": "string",
117 "description": "Full name of the author",
118 },
119 "num_pages": {
120 "type": "integer",
121 "description": "Number of pages in the book",
122 },
123 "short_summary": {
124 "type": "string",
125 "description": "A one-sentence summary of the book",
126 },
127 "long_summary": {
128 "type": "string",
129 "description": (
130 "A detailed two to three sentence summary covering "
131 "the main themes and plot"
132 ),
133 },
134 },
135 "required": ["author", "num_pages", "short_summary", "long_summary"],
136 }
137 system_msg = {
138 "role": "system",
139 "content": (
140 "You are a literary analyst. Extract structured information "
141 "from book descriptions."
142 ),
143 }
144 data = post_batch(
145 {
146 "model": MODEL,
147 "messages": [
148 [
149 system_msg,
150 {
151 "role": "user",
152 "content": (
153 "Extract information from this book: '1984' by George"
154 " Orwell, published in 1949, 328 pages. A dystopian"
155 " novel set in a totalitarian society ruled by Big"
156 " Brother, following Winston Smith as he secretly"
157 " rebels against the oppressive Party that surveils"
158 " and controls every aspect of life."
159 ),
160 },
161 ],
162 [
163 system_msg,
164 {
165 "role": "user",
166 "content": (
167 "Extract information from this book: 'The Hitchhiker's"
168 " Guide to the Galaxy' by Douglas Adams, published in"
169 " 1979, 193 pages. A comedic science fiction novel"
170 " following Arthur Dent, an ordinary Englishman who is"
171 " whisked off Earth moments before it is demolished to"
172 " make way for a hyperspace bypass, and his subsequent"
173 " absurd adventures across the universe."
174 ),
175 },
176 ],
177 ],
178 "response_format": {
179 "type": "json_schema",
180 "json_schema": {
181 "name": "book_summary",
182 "strict": True,
183 "schema": book_schema,
184 },
185 },
186 }
187 )
188 for choice in data["choices"]:
189 book = json.loads(choice["message"]["content"])
190 print(f" [{choice['index']}] {book}")
191
192
193 if __name__ == "__main__":
194 main()