Update messagers/message_outputer.py

messagers/message_outputer.py

@@ -8,10 +8,10 @@ class OpenaiStreamOutputer:
 
     def __init__(self):
         self.default_data = {
-            "created": 1700000000,
-            "id": "chatcmpl-hugginface",
+            "id": "chatcmpl-123",
             "object": "chat.completion",
-            "model": "hugginface",
+            "created": 1677652288,
+            "model": "gpt-3.5-turbo-0613",
             "system_fingerprint": "fp_44709d6fcb",
             "choices": [],
             "usage": {
@@ -21,25 +21,71 @@ class OpenaiStreamOutputer:
             }
         }
 
-    def data_to_string(self, data={}):
+    def data_to_string(self, data={}, content_type=""):
         data_str = f"{json.dumps(data)}"
         return data_str
 
-    def output(self, content=None, role="assistant", content_type="Completions", tokens_info=None) -> str:
+    def output(self, content=None, content_type="Completions", tokens_count=0) -> str:
         data = self.default_data.copy()
-        message = {"role": role, "content": content}
+        if content_type == "Role":
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
+                    "finish_reason": "stop"
+                }
+            ]
+        elif content_type in [
+            "Completions",
+            "InternalSearchQuery",
+            "InternalSearchResult",
+            "SuggestedResponses",
+        ]:
+            if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
+                content += "\n"
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "user",
+                        "content": content,
+                    },
+                    "logprobs": None,
+                    "finish_reason": None,
+                }
+            ]
+        elif content_type == "Finished":
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
+                    "finish_reason": "stop",
+                }
+            ]
+        else:
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": content,
+                    },
+                    "logprobs": None,
+                    "finish_reason": None,
+                }
+            ]
 
-        if tokens_info is None:
-            tokens_info = {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}
+        # 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"]
 
-        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)
+        return self.data_to_string(data, content_type)