fixing streaming

packages/node/meshnet_node/torch_server.py

@@ -342,6 +342,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
         generated: list[str] = []
         current_text = prompt_text
 
+        stream_emit = None
+        if stream:
+            stream_emit = self._start_openai_stream(model_name)
+
         for _ in range(max_tokens):
             try:
                 payload = backend.encode_prompt(current_text)
@@ -357,9 +361,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             if eos_token and token_str == eos_token:
                 break
             generated.append(token_str)
+            if stream_emit is not None:
+                stream_emit(token_str)
             current_text = current_text + token_str
 
         result_text = "".join(generated)
+        if stream_emit is not None:
+            stream_emit(None)
+            return
         self._send_openai_response(result_text, model_name, stream, messages)
 
     def _get_remaining_route(self, model: str) -> list[dict]:
@@ -526,6 +535,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
 
     def _stream_openai_response(self, token_iter, model: str) -> None:
         """Stream tokens from an iterator as SSE chunks."""
+        emit = self._start_openai_stream(model)
+        for token_text in token_iter:
+            if not token_text:
+                continue
+            emit(token_text)
+        emit(None)
+
+    def _start_openai_stream(self, model: str):
+        """Open an OpenAI-compatible SSE response and return a token emitter."""
         chunk_id = "chatcmpl-node"
         created = int(time.time())
         self.send_response(200)
@@ -537,7 +555,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             try:
                 self.wfile.write(f"data: {data}\n\n".encode())
                 self.wfile.flush()
-            except BrokenPipeError:
+            except (BrokenPipeError, ConnectionResetError):
                 pass
 
         _emit(json.dumps({
@@ -545,24 +563,27 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             "model": model,
             "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
         }))
-        for token_text in token_iter:
-            if not token_text:
-                continue
+
+        def emit_token(token_text: str | None) -> None:
+            if token_text is None:
+                _emit(json.dumps({
+                    "id": chunk_id, "object": "chat.completion.chunk", "created": created,
+                    "model": model,
+                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+                }))
+                try:
+                    self.wfile.write(b"data: [DONE]\n\n")
+                    self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError):
+                    pass
+                return
             _emit(json.dumps({
                 "id": chunk_id, "object": "chat.completion.chunk", "created": created,
                 "model": model,
                 "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
             }))
-        _emit(json.dumps({
-            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
-            "model": model,
-            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
-        }))
-        try:
-            self.wfile.write(b"data: [DONE]\n\n")
-            self.wfile.flush()
-        except BrokenPipeError:
-            pass
+
+        return emit_token
 
     def _send_openai_response(
         self,