[PD] add test for chat completions endpoint

Signed-off-by: Abirdcfly <[email protected]>

Author: Abirdcfly

tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py

@@ -52,19 +52,23 @@ def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
 
 
 def run_simple_prompt(base_url: str, model_name: str,
-                      input_prompt: str) -> str:
+                      input_prompt: str, use_chat_endpoint: bool) -> str:
     client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
-    completion = client.completions.create(model=model_name,
-                                           prompt=input_prompt,
-                                           max_tokens=MAX_OUTPUT_LEN,
-                                           temperature=0.0,
-                                           seed=42)
+    if use_chat_endpoint:
+        completion = client.chat.completions.create(model=model_name,
+                                                    messages=[{"role": "user", "content": [{"type": "text","text": input_prompt}]}],
+                                                    max_completion_tokens=MAX_OUTPUT_LEN,
+                                                    temperature=0.0,
+                                                    seed=42)
+        return completion.choices[0].message.content
+    else:
+        completion = client.completions.create(model=model_name,
+                                               prompt=input_prompt,
+                                               max_tokens=MAX_OUTPUT_LEN,
+                                               temperature=0.0,
+                                               seed=42)
 
-    # print("-" * 50)
-    # print(f"Completion results for {model_name}:")
-    # print(completion)
-    # print("-" * 50)
-    return completion.choices[0].text
+        return completion.choices[0].text
 
 
 def main():
@@ -125,10 +129,12 @@ def main():
             f"vllm server: {args.service_url} is not ready yet!")
 
     output_strs = dict()
-    for prompt in SAMPLE_PROMPTS:
+    for i, prompt in enumerate(SAMPLE_PROMPTS):
+        use_chat_endpoint = (i % 2 == 1)
         output_str = run_simple_prompt(base_url=service_url,
                                        model_name=args.model_name,
-                                       input_prompt=prompt)
+                                       input_prompt=prompt,
+                                       use_chat_endpoint=use_chat_endpoint)
         print(f"Prompt: {prompt}, output: {output_str}")
         output_strs[prompt] = output_str
 

tests/v1/kv_connector/nixl_integration/toy_proxy_server.py

@@ -162,6 +162,8 @@ async def send_request_to_service(client_info: dict, endpoint: str,
     }
     req_data["stream"] = False
     req_data["max_tokens"] = 1
+    if "max_completion_tokens" in req_data:
+        req_data["max_completion_tokens"] = 1
     if "stream_options" in req_data:
         del req_data["stream_options"]
     headers = {