Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -1,65 +1,44 @@
 import json
-
+import time
 
 class OpenaiStreamOutputer:
-    """
-    Create chat completion - OpenAI API Documentation
-    * https://platform.openai.com/docs/api-reference/chat/create
-    """
-
     def __init__(self):
         self.default_data = {
-            "created": 1700000000,
-            "id": "chatcmpl-hugginface",
-            "object": "chat.completion.chunk",
-            # "content_type": "Completions",
-            "model": "hugginface",
+            "id": "chatcmpl-123",
+            "object": "chat.completion",
+            "created": None,  # Will be set to current time in output
+            "model": "gpt-3.5-turbo-0613",
+            "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 data_to_string(self, data):
+        return json.dumps(data, indent=2)
 
-    def output(self, content=None, content_type="Completions") -> str:
+    def output(self, content=None, role="assistant", prompt_tokens=0, completion_tokens=0):
         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)
+
+        # Set the current Unix timestamp
+        data["created"] = int(time.time())
+
+        data["choices"] = [{
+            "index": 0,
+            "message": {
+                "role": role,
+                "content": content,
+            },
+            "logprobs": None,
+            "finish_reason": "stop"
+        }]
+
+        # Update token usage
+        data["usage"]["prompt_tokens"] = prompt_tokens
+        data["usage"]["completion_tokens"] = completion_tokens
+        data["usage"]["total_tokens"] = prompt_tokens + completion_tokens
+
+        return self.data_to_string(data)