diff --git a/examples/streaming/test.py b/examples/streaming/test.py index 8041395..3c1b0ec 100644 --- 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 index 45e1f24..39966e8 100644 --- 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 index ed94f1b..d1a121b 100644 --- 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