fix: decoding error when chunking messages

langgenius · Nov 15, 2024 · 9bddcd4 · 9bddcd4
1 parent c1d8441
commit 9bddcd4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/python/dify_plugin/core/server/stdio/request_reader.py b/python/dify_plugin/core/server/stdio/request_reader.py
@@ -14,24 +14,25 @@
 
 class StdioRequestReader(RequestReader):
     def _read_stream(self) -> Generator[PluginInStream, None, None]:
-        buffer = ""
+        buffer = b""
         while True:
             # read data from stdin through tp_read
-            data = tp_read(sys.stdin.fileno(), 512).decode()
+            data = tp_read(sys.stdin.fileno(), 512)
 
             if not data:
                 continue
+
             buffer += data
 
             # process line by line and keep the last line if it is not complete
-            lines = buffer.split("\n")
+            lines = buffer.split(b"\n")
             if len(lines) == 0:
                 continue
 
-            if lines[-1] != "":
+            if lines[-1] != b"":
                 buffer = lines[-1]
             else:
-                buffer = ""
+                buffer = b""
 
             lines = lines[:-1]
             for line in lines:

diff --git a/python/dify_plugin/core/server/tcp/request_reader.py b/python/dify_plugin/core/server/tcp/request_reader.py
@@ -108,14 +108,14 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
         """
         Read data from the target
         """
-        buffer = ""
+        buffer = b""
         while self.alive:
             try:
                 ready_to_read, _, _ = select([self.sock], [], [], 1)
                 if not ready_to_read:
                     continue
-                data = self.sock.recv(4096).decode()
-                if data == "":
+                data = self.sock.recv(4096)
+                if data == b"":
                     raise Exception("Connection is closed")
             except Exception as e:
                 logger.error(f"\033[31mFailed to read data from {self.host}:{self.port}, {e}\033[0m")
@@ -130,14 +130,14 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
             buffer += data
 
             # process line by line and keep the last line if it is not complete
-            lines = buffer.split("\n")
+            lines = buffer.split(b"\n")
             if len(lines) == 0:
                 continue
 
             if lines[-1] != "":
                 buffer = lines[-1]
             else:
-                buffer = ""
+                buffer = b""
 
             lines = lines[:-1]
             for line in lines: