refac

backend/open_webui/utils/misc.py

@@ -106,7 +106,7 @@ def openai_chat_message_template(model: str):
 
 
 def openai_chat_chunk_message_template(
-    model: str, message: Optional[str] = None
+    model: str, message: Optional[str] = None, usage: Optional[dict] = None
 ) -> dict:
     template = openai_chat_message_template(model)
     template["object"] = "chat.completion.chunk"
@@ -114,17 +114,23 @@ def openai_chat_chunk_message_template(
         template["choices"][0]["delta"] = {"content": message}
     else:
         template["choices"][0]["finish_reason"] = "stop"
+
+    if usage:
+        template["usage"] = usage
     return template
 
 
 def openai_chat_completion_message_template(
-    model: str, message: Optional[str] = None
+    model: str, message: Optional[str] = None, usage: Optional[dict] = None
 ) -> dict:
     template = openai_chat_message_template(model)
     template["object"] = "chat.completion"
     if message is not None:
         template["choices"][0]["message"] = {"content": message, "role": "assistant"}
     template["choices"][0]["finish_reason"] = "stop"
+
+    if usage:
+        template["usage"] = usage
     return template