Support v1/chat/completions (#50)

test/srt/test_openai_server.py

@@ -1,8 +1,16 @@
 """
-python3 -m sglang.launch_server --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --port 30000
+First run the following command to launch the server.
+Note that TinyLlama adopts different chat templates in different versions.
+For v0.4, the chat template is chatml.
 
-Output:
-The capital of France is Paris.\nThe capital of the United States is Washington, D.C.\nThe capital of Canada is Ottawa.\nThe capital of Japan is Tokyo
+python3 -m sglang.launch_server --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 \
+--port 30000 --chat-template chatml
+
+Output example:
+The capital of France is Paris.
+The capital of the United States is Washington, D.C.
+The capital of Canada is Ottawa.
+The capital of Japan is Tokyo
 """
 
 import argparse
@@ -38,13 +46,57 @@ def test_completion_stream(args):
     for r in response:
         print(r.choices[0].text, end="", flush=True)
         assert r.id
-        assert r.created
         assert r.usage.prompt_tokens > 0
         assert r.usage.completion_tokens > 0
         assert r.usage.total_tokens > 0
     print()
 
 
+def test_chat_completion(args):
+    client = openai.Client(api_key="EMPTY", base_url=args.base_url)
+    response = client.chat.completions.create(
+        model="default",
+        messages=[
+            {"role": "system", "content": "You are a helpful AI assistant"},
+            {"role": "user", "content": "What is the capital of France?"},
+        ],
+        temperature=0,
+        max_tokens=32,
+    )
+    print(response.choices[0].message.content)
+    assert response.id
+    assert response.created
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.completion_tokens > 0
+    assert response.usage.total_tokens > 0
+
+
+def test_chat_completion_stream(args):
+    client = openai.Client(api_key="EMPTY", base_url=args.base_url)
+    response = client.chat.completions.create(
+        model="default",
+        messages=[
+            {"role": "system", "content": "You are a helpful AI assistant"},
+            {"role": "user", "content": "List 3 countries and their capitals."},
+        ],
+        temperature=0,
+        max_tokens=64,
+        stream=True,
+    )
+    is_first = True
+    for chunk in response:
+        if is_first:
+            is_first = False
+            assert chunk.choices[0].delta.role == "assistant"
+            continue
+
+        data = chunk.choices[0].delta
+        if not data.content:
+            continue
+        print(data.content, end="", flush=True)
+    print()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--base-url", type=str, default="http://127.0.0.1:30000/v1")
@@ -52,3 +104,5 @@ if __name__ == "__main__":
 
     test_completion(args)
     test_completion_stream(args)
+    test_chat_completion(args)
+    test_chat_completion_stream(args)