From 2e0fa1c6a09fbf1ba35c5286458838a96f6d8d1a Mon Sep 17 00:00:00 2001
From: Michael Poluektov <michael.poluektov@gmail.com>
Date: Wed, 31 Jul 2024 22:00:00 +0100
Subject: [PATCH] refac: rename stream_message_template

---
 backend/apps/webui/main.py | 12 ++++++++----
 backend/utils/misc.py      |  2 +-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py
index 46fc00af3..69493f5cc 100644
--- a/backend/apps/webui/main.py
+++ b/backend/apps/webui/main.py
@@ -20,7 +20,7 @@ from apps.webui.models.models import Models
 from apps.webui.utils import load_function_module_by_id
 
 from utils.misc import (
-    stream_message_template,
+    openai_chat_chunk_message_template,
     openai_chat_completion_message_template,
     add_or_update_system_message,
 )
@@ -227,7 +227,7 @@ def process_line(form_data: dict, line):
     if line.startswith("data:"):
         return f"{line}\n\n"
     else:
-        line = stream_message_template(form_data["model"], line)
+        line = openai_chat_chunk_message_template(form_data["model"], line)
         return f"data: {json.dumps(line)}\n\n"
 
 
@@ -371,7 +371,9 @@ async def generate_function_chat_completion(form_data, user):
                     return
 
                 if isinstance(res, str):
-                    message = stream_message_template(form_data["model"], res)
+                    message = openai_chat_chunk_message_template(
+                        form_data["model"], res
+                    )
                     yield f"data: {json.dumps(message)}\n\n"
 
                 if isinstance(res, Iterator):
@@ -383,7 +385,9 @@ async def generate_function_chat_completion(form_data, user):
                         yield process_line(form_data, line)
 
                 if isinstance(res, str) or isinstance(res, Generator):
-                    finish_message = stream_message_template(form_data["model"], "")
+                    finish_message = openai_chat_chunk_message_template(
+                        form_data["model"], ""
+                    )
                     finish_message["choices"][0]["finish_reason"] = "stop"
                     yield f"data: {json.dumps(finish_message)}\n\n"
                     yield "data: [DONE]"
diff --git a/backend/utils/misc.py b/backend/utils/misc.py
index 5e8a39af8..adbf4f8b2 100644
--- a/backend/utils/misc.py
+++ b/backend/utils/misc.py
@@ -97,7 +97,7 @@ def message_template(model: str):
     }
 
 
-def stream_message_template(model: str, message: str):
+def openai_chat_chunk_message_template(model: str, message: str):
     template = message_template(model)
     template["object"] = "chat.completion.chunk"
     template["choices"][0]["delta"] = {"content": message}