| 1 | # ruff: noqa: E501 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | import argparse |
| 5 | import asyncio |
| 6 | import enum |
| 7 | import os |
| 8 | from typing import Any, Literal |
| 9 | |
| 10 | import openai |
| 11 | import pydantic |
| 12 | from openai.types.chat import ChatCompletionChunk |
| 13 | |
| 14 | ConstraintsFormat = Literal[ |
| 15 | "choice", |
| 16 | "regex", |
| 17 | "json", |
| 18 | "grammar", |
| 19 | "structural_tag", |
| 20 | ] |
| 21 | |
| 22 | |
| 23 | async def print_stream_response( |
| 24 | stream_response: openai.AsyncStream[ChatCompletionChunk], |
| 25 | title: str, |
| 26 | args: argparse.Namespace, |
| 27 | ): |
| 28 | print(f"\n\n{title} (Streaming):") |
| 29 | |
| 30 | local_reasoning_header_printed = False |
| 31 | local_content_header_printed = False |
| 32 | |
| 33 | async for chunk in stream_response: |
| 34 | delta = chunk.choices[0].delta |
| 35 | |
| 36 | reasoning_chunk_text: str | None = getattr(delta, "reasoning", None) |
| 37 | content_chunk_text = delta.content |
| 38 | |
| 39 | if args.reasoning: |
| 40 | if reasoning_chunk_text: |
| 41 | if not local_reasoning_header_printed: |
| 42 | print(" Reasoning: ", end="") |
| 43 | local_reasoning_header_printed = True |
| 44 | print(reasoning_chunk_text, end="", flush=True) |
| 45 | |
| 46 | if content_chunk_text: |
| 47 | if not local_content_header_printed: |
| 48 | if local_reasoning_header_printed: |
| 49 | print() |
| 50 | print(" Content: ", end="") |
| 51 | local_content_header_printed = True |
| 52 | print(content_chunk_text, end="", flush=True) |
| 53 | else: |
| 54 | if content_chunk_text: |
| 55 | if not local_content_header_printed: |
| 56 | print(" Content: ", end="") |
| 57 | local_content_header_printed = True |
| 58 | print(content_chunk_text, end="", flush=True) |
| 59 | print() |
| 60 | |
| 61 | |
| 62 | class CarType(str, enum.Enum): |
| 63 | SEDAN = "SEDAN" |
| 64 | SUV = "SUV" |
| 65 | TRUCK = "TRUCK" |
| 66 | COUPE = "COUPE" |
| 67 | |
| 68 | |
| 69 | class CarDescription(pydantic.BaseModel): |
| 70 | brand: str |
| 71 | model: str |
| 72 | car_type: CarType |
| 73 | |
| 74 | |
| 75 | PARAMS: dict[ConstraintsFormat, dict[str, Any]] = { |
| 76 | "choice": { |
| 77 | "messages": [ |
| 78 | { |
| 79 | "role": "user", |
| 80 | "content": "Classify this sentiment: vLLM is wonderful!", |
| 81 | } |
| 82 | ], |
| 83 | "extra_body": {"structured_outputs": {"choice": ["positive", "negative"]}}, |
| 84 | }, |
| 85 | "regex": { |
| 86 | "messages": [ |
| 87 | { |
| 88 | "role": "user", |
| 89 | "content": "Generate an email address for Alan Turing, who works in Enigma. End in .com and new line. Example result: 'alan.turing@enigma.com\n'", |
| 90 | } |
| 91 | ], |
| 92 | "extra_body": { |
| 93 | "structured_outputs": {"regex": r"[a-z0-9.]{1,20}@\w{6,10}\.com\n"}, |
| 94 | }, |
| 95 | }, |
| 96 | "json": { |
| 97 | "messages": [ |
| 98 | { |
| 99 | "role": "user", |
| 100 | "content": "Generate a JSON with the brand, model and car_type of the most iconic car from the 90's", |
| 101 | } |
| 102 | ], |
| 103 | "response_format": { |
| 104 | "type": "json_schema", |
| 105 | "json_schema": { |
| 106 | "name": "car-description", |
| 107 | "schema": CarDescription.model_json_schema(), |
| 108 | }, |
| 109 | }, |
| 110 | }, |
| 111 | "grammar": { |
| 112 | "messages": [ |
| 113 | { |
| 114 | "role": "user", |
| 115 | "content": "Generate an SQL query to show the 'username' and 'email' from the 'users' table.", |
| 116 | } |
| 117 | ], |
| 118 | "extra_body": { |
| 119 | "structured_outputs": { |
| 120 | "grammar": """ |
| 121 | root ::= select_statement |
| 122 | |
| 123 | select_statement ::= "SELECT " column " from " table " where " condition |
| 124 | |
| 125 | column ::= "col_1 " | "col_2 " |
| 126 | |
| 127 | table ::= "table_1 " | "table_2 " |
| 128 | |
| 129 | condition ::= column "= " number |
| 130 | |
| 131 | number ::= "1 " | "2 " |
| 132 | """, |
| 133 | } |
| 134 | }, |
| 135 | }, |
| 136 | "structural_tag": { |
| 137 | "messages": [ |
| 138 | { |
| 139 | "role": "user", |
| 140 | "content": """ |
| 141 | You have access to the following function to retrieve the weather in a city: |
| 142 | |
| 143 | { |
| 144 | "name": "get_weather", |
| 145 | "parameters": { |
| 146 | "city": { |
| 147 | "param_type": "string", |
| 148 | "description": "The city to get the weather for", |
| 149 | "required": True |
| 150 | } |
| 151 | } |
| 152 | } |
| 153 | |
| 154 | If a you choose to call a function ONLY reply in the following format: |
| 155 | <{start_tag}={function_name}>{parameters}{end_tag} |
| 156 | where |
| 157 | |
| 158 | start_tag => `<function` |
| 159 | parameters => a JSON dict with the function argument name as key and function |
| 160 | argument value as value. |
| 161 | end_tag => `</function>` |
| 162 | |
| 163 | Here is an example, |
| 164 | <function=example_function_name>{"example_name": "example_value"}</function> |
| 165 | |
| 166 | Reminder: |
| 167 | - Function calls MUST follow the specified format |
| 168 | - Required parameters MUST be specified |
| 169 | - Only call one function at a time |
| 170 | - Put the entire function call reply on one line |
| 171 | - Always add your sources when using search results to answer the user query |
| 172 | |
| 173 | You are a helpful assistant. |
| 174 | |
| 175 | Given the previous instructions, what is the weather in New York City, Boston, |
| 176 | and San Francisco?""", |
| 177 | }, |
| 178 | ], |
| 179 | "response_format": { |
| 180 | "type": "structural_tag", |
| 181 | "structures": [ |
| 182 | { |
| 183 | "begin": "<function=get_weather>", |
| 184 | "schema": { |
| 185 | "type": "object", |
| 186 | "properties": {"city": {"type": "string"}}, |
| 187 | "required": ["city"], |
| 188 | }, |
| 189 | "end": "</function>", |
| 190 | } |
| 191 | ], |
| 192 | "triggers": ["<function="], |
| 193 | }, |
| 194 | }, |
| 195 | } |
| 196 | |
| 197 | |
| 198 | async def cli(): |
| 199 | parser = argparse.ArgumentParser( |
| 200 | description="Run OpenAI Chat Completion with various structured outputs capabilities", |
| 201 | ) |
| 202 | _ = parser.add_argument( |
| 203 | "--constraint", |
| 204 | type=str, |
| 205 | nargs="+", |
| 206 | choices=[*list(PARAMS), "*"], |
| 207 | default=["*"], |
| 208 | help="Specify which constraint(s) to run.", |
| 209 | ) |
| 210 | _ = parser.add_argument( |
| 211 | "--stream", |
| 212 | action=argparse.BooleanOptionalAction, |
| 213 | default=False, |
| 214 | help="Enable streaming output", |
| 215 | ) |
| 216 | _ = parser.add_argument( |
| 217 | "--reasoning", |
| 218 | action=argparse.BooleanOptionalAction, |
| 219 | default=False, |
| 220 | help="Enable printing of reasoning traces if available.", |
| 221 | ) |
| 222 | args = parser.parse_args() |
| 223 | |
| 224 | base_url = os.getenv("OPENAI_BASE_URL", "http://localhost:8000/v1") |
| 225 | client = openai.AsyncOpenAI(base_url=base_url, api_key="EMPTY") |
| 226 | constraints = list(PARAMS) if "*" in args.constraint else list(set(args.constraint)) |
| 227 | model = (await client.models.list()).data[0].id |
| 228 | |
| 229 | if args.stream: |
| 230 | results = await asyncio.gather( |
| 231 | *[ |
| 232 | client.chat.completions.create( |
| 233 | model=model, |
| 234 | max_tokens=1024, |
| 235 | stream=True, |
| 236 | **PARAMS[name], |
| 237 | ) |
| 238 | for name in constraints |
| 239 | ] |
| 240 | ) |
| 241 | for constraint, stream in zip(constraints, results): |
| 242 | await print_stream_response(stream, constraint, args) |
| 243 | else: |
| 244 | results = await asyncio.gather( |
| 245 | *[ |
| 246 | client.chat.completions.create( |
| 247 | model=model, |
| 248 | max_tokens=1024, |
| 249 | stream=False, |
| 250 | **PARAMS[name], |
| 251 | ) |
| 252 | for name in constraints |
| 253 | ] |
| 254 | ) |
| 255 | for constraint, response in zip(constraints, results): |
| 256 | print(f"\n\n{constraint}:") |
| 257 | message = response.choices[0].message |
| 258 | if args.reasoning and hasattr(message, "reasoning"): |
| 259 | print(f" Reasoning: {message.reasoning or ''}") |
| 260 | print(f" Content: {message.content!r}") |
| 261 | |
| 262 | |
| 263 | def main(): |
| 264 | asyncio.run(cli()) |
| 265 | |
| 266 | |
| 267 | if __name__ == "__main__": |
| 268 | main() |