API Client#

Source vllm-project/vllm.

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