Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -1,44 +1,65 @@
 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 = {
-            "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",
+            "created": 1700000000,
+            "id": "chatcmpl-hugginface",
+            "object": "chat.completion.chunk",
+            # "content_type": "Completions",
+            "model": "hugginface",
             "choices": [],
-            "usage": {
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
-                "total_tokens": 0
-            }
         }
 
-    def data_to_string(self, data):
-        return json.dumps(data, indent=2)
+    def data_to_string(self, data={}, content_type=""):
+        data_str = f"{json.dumps(data)}"
+        return data_str
 
-    def output(self, content=None, role="assistant", prompt_tokens=0, completion_tokens=0):
+    def output(self, content=None, content_type="Completions") -> str:
         data = self.default_data.copy()
-
-        # 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)
+        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)