API Client#

Source vllm-project/vllm.

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