From 5da0de417e2b2a53f3b4db1fd6c92211f64b9e26 Mon Sep 17 00:00:00 2001 From: Francisco Castillo Date: Mon, 21 Oct 2024 16:27:37 -0700 Subject: [PATCH] Use Response if only one dataframe, stream otherwise --- src/phoenix/server/api/routers/v1/spans.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/phoenix/server/api/routers/v1/spans.py b/src/phoenix/server/api/routers/v1/spans.py index 52c4d5778a..827b2a487b 100644 --- a/src/phoenix/server/api/routers/v1/spans.py +++ b/src/phoenix/server/api/routers/v1/spans.py @@ -120,14 +120,21 @@ async def query_spans_handler( raise HTTPException(status_code=HTTP_404_NOT_FOUND) print(f"{results=}\n\n") - async def content() -> AsyncIterator[bytes]: - for result in results: - yield df_to_bytes(result) + if len(results) == 1: + return Response( + content=df_to_bytes(results[0]), + media_type="application/x-pandas-arrow", + ) + else: - return Response( - content=content(), - media_type="application/x-pandas-arrow", - ) + async def content() -> AsyncIterator[bytes]: + for result in results: + yield df_to_bytes(result) + + return StreamingResponse( + content=content(), + media_type="application/x-pandas-arrow", + ) @router.get("/spans", include_in_schema=False, deprecated=True)