Use Response if only one dataframe, stream otherwise

Arize-ai · Oct 21, 2024 · 5da0de4 · 5da0de4
1 parent bbeaa01
commit 5da0de4
Showing 1 changed file with 14 additions and 7 deletions.
diff --git 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)