Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -1,6 +1,5 @@
 import json
 
-
 class OpenaiStreamOutputer:
     """
     Create chat completion - OpenAI API Documentation
@@ -11,55 +10,36 @@ class OpenaiStreamOutputer:
         self.default_data = {
             "created": 1700000000,
             "id": "chatcmpl-hugginface",
-            "object": "chat.completion.chunk",
-            # "content_type": "Completions",
+            "object": "chat.completion",
             "model": "hugginface",
+            "system_fingerprint": "fp_44709d6fcb",
             "choices": [],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            }
         }
 
-    def data_to_string(self, data={}, content_type=""):
+    def data_to_string(self, data={}):
         data_str = f"{json.dumps(data)}"
         return data_str
 
-    def output(self, content=None, content_type="Completions") -> str:
+    def output(self, content=None, role="assistant", content_type="Completions", tokens_info=None) -> str:
         data = self.default_data.copy()
-        if content_type == "Role":
-            data["choices"] = [
-                {
-                    "index": 0,
-                    "delta": {"role": "assistant"},
-                    "finish_reason": None,
-                }
-            ]
-        elif content_type in [
-            "Completions",
-            "InternalSearchQuery",
-            "InternalSearchResult",
-            "SuggestedResponses",
-        ]:
-            if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
-                content += "\n"
-            data["choices"] = [
-                {
-                    "index": 0,
-                    "delta": {"content": content},
-                    "finish_reason": None,
-                }
-            ]
-        elif content_type == "Finished":
-            data["choices"] = [
-                {
-                    "index": 0,
-                    "delta": {},
-                    "finish_reason": "stop",
-                }
-            ]
-        else:
-            data["choices"] = [
-                {
-                    "index": 0,
-                    "delta": {},
-                    "finish_reason": None,
-                }
-            ]
-        return self.data_to_string(data, content_type)
+        message = {"role": role, "content": content}
+
+        if tokens_info is None:
+            tokens_info = {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}
+
+        data["choices"] = [
+            {
+                "index": 0,
+                "message": message,
+                "logprobs": None,
+                "finish_reason": "stop" if content_type == "Finished" else None
+            }
+        ]
+
+        data["usage"] = tokens_info
+        return self.data_to_string(data)