From a19ad95724a3fb8d87e9aa73a2b3d519744ab93a Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Wed, 22 Jan 2025 14:37:45 +0800
Subject: [PATCH] fix: too many message exception error #275

---
 src/hugchat/hugchat.py | 21 ++++++++++++++++-----
 src/hugchat/message.py |  3 ++-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/hugchat/hugchat.py b/src/hugchat/hugchat.py
index c828687..2b12416 100644
--- 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
index d762942..b963ad6 100644
--- 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"]