Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -1,6 +1,5 @@
 import json
-import time 
-
+import time
 
 class OpenaiStreamOutputer:
     """
@@ -9,28 +8,37 @@ class OpenaiStreamOutputer:
     """
 
     def __init__(self):
-        current_time = int(time.time()) 
+        current_time = int(time.time())
         self.default_data = {
-            "created": 1700000,
             "id": "chatcmpl-hugginface",
             "object": "chat.completion.chunk",
-            # "content_type": "Completions",
+            "created": current_time,
             "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=""):
         data_str = f"{json.dumps(data)}"
         return data_str
 
-    def output(self, content=None, content_type="Completions") -> str:
+    def output(self, content=None, content_type="Completions", tokens_count=0) -> str:
         data = self.default_data.copy()
         if content_type == "Role":
             data["choices"] = [
                 {
                     "index": 0,
-                    "delta": {"role": "assistant"},
-                    "finish_reason": None,
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
+                    "finish_reason": "stop"
                 }
             ]
         elif content_type in [
@@ -44,7 +52,11 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "delta": {"content": content},
+                    "message": {
+                        "role": "user",
+                        "content": content,
+                    },
+                    "logprobs": None,
                     "finish_reason": None,
                 }
             ]
@@ -52,7 +64,11 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "delta": {},
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
                     "finish_reason": "stop",
                 }
             ]
@@ -60,8 +76,18 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "delta": {},
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
                     "finish_reason": None,
                 }
             ]
-        return self.data_to_string(data, content_type)
+
+        # Update token counts
+        data["usage"]["prompt_tokens"] += tokens_count
+        data["usage"]["completion_tokens"] += len(content.split())
+        data["usage"]["total_tokens"] = data["usage"]["prompt_tokens"] + data["usage"]["completion_tokens"]
+
+        return self.data_to_string(data, content_type)