| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Example Python client for `examples/applications/api_server/server.py` |
| 4 | Start the demo server: |
| 5 | python examples/applications/api_server/server.py --model <model_name> |
| 6 | |
| 7 | NOTE: The API server is used only for demonstration and simple performance |
| 8 | benchmarks. It is not intended for production use. |
| 9 | For production use, we recommend `vllm serve` and the OpenAI client API. |
| 10 | """ |
| 11 | |
| 12 | import argparse |
| 13 | import json |
| 14 | from argparse import Namespace |
| 15 | from collections.abc import Iterable |
| 16 | |
| 17 | import requests |
| 18 | |
| 19 | |
| 20 | def clear_line(n: int = 1) -> None: |
| 21 | LINE_UP = "\033[1A" |
| 22 | LINE_CLEAR = "\x1b[2K" |
| 23 | for _ in range(n): |
| 24 | print(LINE_UP, end=LINE_CLEAR, flush=True) |
| 25 | |
| 26 | |
| 27 | def post_http_request( |
| 28 | prompt: str, api_url: str, n: int = 1, stream: bool = False |
| 29 | ) -> requests.Response: |
| 30 | headers = {"User-Agent": "Test Client"} |
| 31 | pload = { |
| 32 | "prompt": prompt, |
| 33 | "n": n, |
| 34 | "temperature": 0.0, |
| 35 | "max_tokens": 16, |
| 36 | "stream": stream, |
| 37 | } |
| 38 | response = requests.post(api_url, headers=headers, json=pload, stream=stream) |
| 39 | return response |
| 40 | |
| 41 | |
| 42 | def get_streaming_response(response: requests.Response) -> Iterable[list[str]]: |
| 43 | for chunk in response.iter_lines( |
| 44 | chunk_size=8192, decode_unicode=False, delimiter=b"\n" |
| 45 | ): |
| 46 | if chunk: |
| 47 | data = json.loads(chunk.decode("utf-8")) |
| 48 | output = data["text"] |
| 49 | yield output |
| 50 | |
| 51 | |
| 52 | def get_response(response: requests.Response) -> list[str]: |
| 53 | data = json.loads(response.content) |
| 54 | output = data["text"] |
| 55 | return output |
| 56 | |
| 57 | |
| 58 | def parse_args(): |
| 59 | parser = argparse.ArgumentParser() |
| 60 | parser.add_argument("--host", type=str, default="localhost") |
| 61 | parser.add_argument("--port", type=int, default=8000) |
| 62 | parser.add_argument("--n", type=int, default=1) |
| 63 | parser.add_argument("--prompt", type=str, default="San Francisco is a") |
| 64 | parser.add_argument("--stream", action="store_true") |
| 65 | return parser.parse_args() |
| 66 | |
| 67 | |
| 68 | def main(args: Namespace): |
| 69 | prompt = args.prompt |
| 70 | api_url = f"http://{args.host}:{args.port}/generate" |
| 71 | n = args.n |
| 72 | stream = args.stream |
| 73 | |
| 74 | print(f"Prompt: {prompt!r}\n", flush=True) |
| 75 | response = post_http_request(prompt, api_url, n, stream) |
| 76 | |
| 77 | if stream: |
| 78 | num_printed_lines = 0 |
| 79 | for h in get_streaming_response(response): |
| 80 | clear_line(num_printed_lines) |
| 81 | num_printed_lines = 0 |
| 82 | for i, line in enumerate(h): |
| 83 | num_printed_lines += 1 |
| 84 | print(f"Beam candidate {i}: {line!r}", flush=True) |
| 85 | else: |
| 86 | output = get_response(response) |
| 87 | for i, line in enumerate(output): |
| 88 | print(f"Beam candidate {i}: {line!r}", flush=True) |
| 89 | |
| 90 | |
| 91 | if __name__ == "__main__": |
| 92 | args = parse_args() |
| 93 | main(args) |