Update app.py

app.py

@@ -14,7 +14,6 @@ CHAT_URL = BASE_URL + "/conversation"
 # Initialize the token and message history
 token = ""
 messHistory: list = []
-
 async def chat(messList):
     """Async function to send and receive messages with the server."""
     global token
@@ -31,29 +30,39 @@ async def chat(messList):
             "stream": True
         }
 
+        fullmessage = ""
         # Make the POST request to the chat API
         async with session.post(CHAT_URL, json=body) as resp:
             if resp.status != 200:
                 return "Error occurred during the chat process."
 
-            fullmessage = ""
-            # Stream and decode response chunks
+            # Use a buffer to accumulate data chunks
+            buffer = ""
             async for chunk in resp.content.iter_any():
-                data_str = chunk.decode("utf-8")
-                json_str = data_str.replace("data: ", "")
-                if json_str.strip() == "[DONE]":
-                    break
-                else:
+                buffer += chunk.decode("utf-8")
+
+                # Attempt to decode JSON objects from the buffer
+                while True:
                     try:
+                        # Find the end of the JSON object
+                        index = buffer.index('\n')
+                        json_str = buffer[:index].strip()
+                        buffer = buffer[index+1:]
+
+                        if json_str.strip() == "[DONE]":
+                            break
+
                         data_dict = json.loads(json_str)
-                        fullmessage += data_dict["message"]
-                        token = data_dict["resp_token"]  # Update token
-                    except KeyError:
-                        pass
+                        fullmessage += data_dict.get("message", "")
+                        token = data_dict.get("resp_token", token)  # Update token
+                    except (ValueError, json.JSONDecodeError):
+                        # Handle incomplete or malformed JSON
+                        break
             
             messHistory.append({"role": "assistant", "content": fullmessage})  # Append assistant response
             return fullmessage
 
+
 def gradio_chat(user_input, mode):
     """Synchronous wrapper for the async chat function, integrated with Gradio."""
     messHistory.append({"role": "user", "content": f"[{mode}] {user_input}"})