Update networks/message_streamer.py

networks/message_streamer.py

@@ -26,24 +26,15 @@ class MessageStreamer:
             self.model = model
         else:
             self.model = "default"
-            self.model_fullname = self.MODEL_MAP[self.model]
-            self.message_outputer = OpenaiStreamOutputer()
-            self.stop_sequences = self.STOP_SEQUENCES_MAP.get(self.model, "")
+        self.model_fullname = self.MODEL_MAP[self.model]
+        self.message_outputer = OpenaiStreamOutputer()
 
     def parse_line(self, line):
-        try:
-            line = line.decode("utf-8")
-            line = re.sub(r"data:\s*", "", line)
-            data = json.loads(line)
-    
-            # Check if 'token' key exists and return its 'text' value
-            if 'token' in data:
-                return data['token'].get('text', '')  # Return an empty string if 'text' is not found
-            else:
-                return None  # Return None if 'token' key is not present
-        except Exception as e:
-            print(f"An error occurred: {e}")
-            return None
+        line = line.decode("utf-8")
+        line = re.sub(r"data:\s*", "", line)
+        data = json.loads(line)
+        content = data["token"]["text"]
+        return content
 
     def chat_response(
         self,
@@ -107,6 +98,7 @@ class MessageStreamer:
         return stream_response
 
     def chat_return_dict(self, stream_response):
+        # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
         final_output = self.message_outputer.default_data.copy()
         final_output["choices"] = [
             {
@@ -118,28 +110,24 @@ class MessageStreamer:
                 },
             }
         ]
-    
         logger.back(final_output)
-    
+
         final_content = ""
         for line in stream_response.iter_lines():
             if not line:
                 continue
-    
             content = self.parse_line(line)
-            
-            # Check if content is None before calling strip()
-            if content is not None:
-                if content.strip() == self.stop_sequences:
-                    logger.success("\n[Finished]")
-                    break
-                else:
-                    logger.back(content, end="")
-                    final_content += content
-    
+
+            if content.strip() == self.stop_sequences:
+                logger.success("\n[Finished]")
+                break
+            else:
+                logger.back(content, end="")
+                final_content += content
+
         if self.model in self.STOP_SEQUENCES_MAP.keys():
             final_content = final_content.replace(self.stop_sequences, "")
-    
+
         final_output["choices"][0]["message"]["content"] = final_content
         return final_output
 
@@ -148,24 +136,21 @@ class MessageStreamer:
         for line in stream_response.iter_lines():
             if not line:
                 continue
-    
+
             content = self.parse_line(line)
 
-            stop_sequence = getattr(self, 'stop_sequences', '')  # Default value if not set
-    
-            if content is not None and content.strip() == self.stop_sequences:
+            if content.strip() == self.stop_sequences:
                 content_type = "Finished"
                 logger.success("\n[Finished]")
                 is_finished = True
             else:
                 content_type = "Completions"
                 logger.back(content, end="")
-    
+
             output = self.message_outputer.output(
                 content=content, content_type=content_type
             )
             yield output
-    
-        if not is_finished:
-            yield self.message_outputer.output(content="", content_type="Finished")
 
+        if not is_finished:
+            yield self.message_outputer.output(content="", content_type="Finished")