API Client

Source vllm-project/vllm.

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