fix: too many message exception error #275

Soulter · Jan 22, 2025 · a19ad95 · a19ad95
1 parent 9d4b863
commit a19ad95
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/src/hugchat/hugchat.py b/src/hugchat/hugchat.py
@@ -3,6 +3,7 @@
 import os
 import datetime
 import logging
+import time
 import typing
 import traceback
 
@@ -698,9 +699,9 @@ def _stream_query(
             'sec-fetch-site': 'same-origin',
             'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0',
         }
-        final_answer = {}
-
+        obj = {}
         break_flag = False
+        has_started = False
 
         while retry_count > 0:
             resp = self.session.post(
@@ -725,13 +726,23 @@ def _stream_query(
                         continue
                     res = line
                     obj = json.loads(res)
-                    if obj.__contains__("type"):
+                    if "type" in obj:
                         _type = obj["type"]
 
                         if _type == "finalAnswer":
-                            final_answer = obj
                             break_flag = True
                             break
+
+                        if _type == "status" and obj["status"] == "started":
+                            if has_started:
+                                obj = {
+                                    "type": "finalAnswer",
+                                    "text": ""
+                                }
+                                break_flag = True
+                                break
+                            has_started = True
+
                     else:
                         logging.error(f"No `type` found in response: {obj}")
                     yield obj
@@ -752,7 +763,7 @@ def _stream_query(
 
         # update the history of current conversation
         self.get_conversation_info(conversation)
-        yield final_answer
+        yield obj
 
     def query(self) -> Message:
         """

diff --git a/src/hugchat/message.py b/src/hugchat/message.py
@@ -127,7 +127,7 @@ def __next__(self) -> dict:
 
             # set _result_text if this is the final iteration of the chat message
             if data_type == ResponseTypes.FINAL:
-                self._result_text = data["text"]
+                # self._result_text = data["text"]
                 self.msg_status = MessageStatus.RESOLVED
 
             # Handle web response type
@@ -156,6 +156,7 @@ def __next__(self) -> dict:
             # replace null characters with an empty string
             elif data_type == ResponseTypes.STREAM:
                 data["token"] = data["token"].replace('\u0000', '')
+                self._result_text += data["token"]
 
             elif "messageType" in data:
                 message_type: str = data["messageType"]