Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -24,22 +24,18 @@ class OpenaiStreamOutputer:
             }
         }
 
-    def data_to_string(self, data={}, content_type=""):
+  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", tokens_count=0) -> str:
+    def output(self, content=None, content_type="Completions") -> str:
         data = self.default_data.copy()
         if content_type == "Role":
             data["choices"] = [
                 {
                     "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": content,
-                    },
-                    "logprobs": None,
-                    "finish_reason": "stop"
+                    "delta": {"role": "assistant"},
+                    "finish_reason": None,
                 }
             ]
         elif content_type in [
@@ -53,11 +49,7 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "message": {
-                        "role": "user",
-                        "content": content,
-                    },
-                    "logprobs": None,
+                    "delta": {"content": content},
                     "finish_reason": None,
                 }
             ]
@@ -65,11 +57,7 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": content,
-                    },
-                    "logprobs": None,
+                    "delta": {},
                     "finish_reason": "stop",
                 }
             ]
@@ -77,20 +65,8 @@ class OpenaiStreamOutputer:
             data["choices"] = [
                 {
                     "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": content,
-                    },
-                    "logprobs": None,
+                    "delta": {},
                     "finish_reason": None,
                 }
             ]
-
-        # Update token counts
-        tokens_used = len(content.split())
-        self.token_manager.add_tokens(tokens_used)  # hypothetical method
-        data["usage"]["prompt_tokens"] += tokens_count
-        data["usage"]["completion_tokens"] += tokens_used
-        data["usage"]["total_tokens"] = self.token_manager.get_total_tokens()  # hypothetical method
-
-        return self.data_to_string(data, content_type)
+        return self.data_to_string(data, content_type)