Support streaming APIs

yangwenz · Jul 5, 2024 · ec03428 · ec03428
1 parent a1df35c
commit ec03428
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/examples/streaming/test.py b/examples/streaming/test.py
@@ -2,5 +2,5 @@
 
 url = "http://localhost:8080/v1/models/streaming:generate"
 with requests.post(url, stream=True, json={"repeat": 5}) as r:
-    for chunk in r.iter_content(16):
+    for chunk in r.iter_lines():
         print(chunk)
diff --git a/kservehelper/kserve/rest/v1_endpoints.py b/kservehelper/kserve/rest/v1_endpoints.py
@@ -78,7 +78,7 @@ async def generate(self, model_name: str, request: Request) -> StreamingResponse
         headers = dict(request.headers.items())
         results_generator, response_headers = \
             await self.dataplane.generate(model_name=model_name, body=body, headers=headers)
-        return StreamingResponse(results_generator())
+        return StreamingResponse(results_generator(), media_type="application/x-ndjson")
 
     async def explain(self, model_name: str, request: Request) -> Union[Response, Dict]:
         """Explain handler.

diff --git a/kservehelper/model.py b/kservehelper/model.py
@@ -310,7 +310,7 @@ def generate_filepath(filename: str) -> str:
     def wrap_generator(g):
         def _g():
             for i, data in enumerate(g()):
-                yield json.dumps({"id": i, "data": data})
+                yield json.dumps({"id": i, "data": data}) + "\n"
 
         return _g