bugfix of openai_api_server for fastchat.serve.vllm_worker (#2398)

Co-authored-by: wuyongyu <wuyongyu@atomecho.xyz>
lm-sys · Sep 11, 2023 · b921f16 · b921f16
1 parent 7dcdafe
commit b921f16
Showing 1 changed file with 40 additions and 14 deletions.
diff --git a/fastchat/serve/openai_api_server.py b/fastchat/serve/openai_api_server.py
@@ -399,15 +399,28 @@ async def create_chat_completion(request: ChatCompletionRequest):
     for i, content in enumerate(all_tasks):
         if content["error_code"] != 0:
             return create_error_response(content["error_code"], content["text"])
-        choices.append(
-            ChatCompletionResponseChoice(
-                index=i,
-                message=ChatMessage(role="assistant", content=content["text"]),
-                finish_reason=content.get("finish_reason", "stop"),
+        if isinstance(content["text"], list):
+            for t in content["text"]:
+                choices.append(
+                    ChatCompletionResponseChoice(
+                        index=i,
+                        message=ChatMessage(role="assistant", content=t),
+                        finish_reason=content.get("finish_reason", "stop"),
+                    )
+                )
+        else:
+            choices.append(
+                ChatCompletionResponseChoice(
+                    index=i,
+                    message=ChatMessage(role="assistant", content=content["text"]),
+                    finish_reason=content.get("finish_reason", "stop"),
+                )
             )
-        )
         if "usage" in content:
-            task_usage = UsageInfo.parse_obj(content["usage"])
+            if isinstance(content["usage"], list):
+                task_usage = UsageInfo.parse_obj(content["usage"][0])
+            else:
+                task_usage = UsageInfo.parse_obj(content["usage"])
             for usage_key, usage_value in task_usage.dict().items():
                 setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
 
@@ -775,14 +788,27 @@ async def create_chat_completion(request: APIChatCompletionRequest):
     for i, content in enumerate(all_tasks):
         if content["error_code"] != 0:
             return create_error_response(content["error_code"], content["text"])
-        choices.append(
-            ChatCompletionResponseChoice(
-                index=i,
-                message=ChatMessage(role="assistant", content=content["text"]),
-                finish_reason=content.get("finish_reason", "stop"),
+        if isinstance(content["text"], list):
+            for t in content["text"]:
+                choices.append(
+                    ChatCompletionResponseChoice(
+                        index=i,
+                        message=ChatMessage(role="assistant", content=t),
+                        finish_reason=content.get("finish_reason", "stop"),
+                    )
+                )
+        else:
+            choices.append(
+                ChatCompletionResponseChoice(
+                    index=i,
+                    message=ChatMessage(role="assistant", content=content["text"]),
+                    finish_reason=content.get("finish_reason", "stop"),
+                )
             )
-        )
-        task_usage = UsageInfo.parse_obj(content["usage"])
+        if isinstance(content["usage"], list):
+            task_usage = UsageInfo.parse_obj(content["usage"][0])
+        else:
+            task_usage = UsageInfo.parse_obj(content["usage"])
         for usage_key, usage_value in task_usage.dict().items():
             setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)