Add role for DeltaMessage

xusenlinzy · Sep 7, 2023 · 7f9c70a · 7f9c70a
1 parent 97d725a
commit 7f9c70a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/api/routes/chat.py b/api/routes/chat.py
@@ -180,7 +180,7 @@ async def chat_completion_stream_generator(
                         found_action_name = True
                         finish_reason = "function_call"
             else:
-                messages = [DeltaMessage(content=delta_text)]
+                messages = [DeltaMessage(content=delta_text, role=Role.ASSISTANT)]
                 finish_reason = content.get("finish_reason", "stop")
 
             chunks = []

diff --git a/api/vllm_routes/chat.py b/api/vllm_routes/chat.py
@@ -189,7 +189,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
                             found_action_name = True
                             finish_reason = "function_call"
                 else:
-                    msgs = [DeltaMessage(content=delta_text)]
+                    msgs = [DeltaMessage(content=delta_text, role=Role.ASSISTANT)]
                     finish_reason = output.finish_reason
 
                 for m in msgs:
@@ -199,7 +199,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
                 if output.finish_reason is not None:
                     response_json = create_stream_response_json(
                         index=i,
-                        delta=DeltaMessage(content=""),
+                        delta=DeltaMessage(content="", role=Role.ASSISTANT),
                         finish_reason=output.finish_reason,
                     )
                     yield f"data: {response_json}\n\n"