Allow downloading run sheet

Author: ajhai

llmstack/client/src/components/sheets/Sheet.jsx

@@ -22,6 +22,7 @@ import { enqueueSnackbar } from "notistack";
 import SaveIcon from "@mui/icons-material/Save";
 import PlayArrowIcon from "@mui/icons-material/PlayArrow";
 import PauseIcon from "@mui/icons-material/Pause";
+import DownloadIcon from "@mui/icons-material/Download";
 
 import "@glideapps/glide-data-grid/dist/index.css";
 
@@ -51,7 +52,14 @@ const gridCellToCellId = (gridCell, columns) => {
   return `${colLetter}${rowIndex + 1}`;
 };
 
-const SheetHeader = ({ sheet, setRunId, hasChanges, onSave, sheetRunning }) => {
+const SheetHeader = ({
+  sheet,
+  setRunId,
+  hasChanges,
+  onSave,
+  sheetRunning,
+  runId,
+}) => {
   const navigate = useNavigate();
 
   const saveSheet = () => {
@@ -83,6 +91,10 @@ const SheetHeader = ({ sheet, setRunId, hasChanges, onSave, sheetRunning }) => {
     }
   };
 
+  const downloadSheet = () => {
+    window.open(`/api/sheets/${sheet.uuid}/download`, "_blank");
+  };
+
   return (
     <Stack>
       <Typography variant="h5" className="section-header">
@@ -124,6 +136,18 @@ const SheetHeader = ({ sheet, setRunId, hasChanges, onSave, sheetRunning }) => {
                 <SaveIcon />
               </Button>
             </Tooltip>
+            {!sheetRunning && (
+              <Tooltip title="Download CSV">
+                <Button
+                  onClick={downloadSheet}
+                  color="primary"
+                  variant="outlined"
+                  sx={{ minWidth: "40px", padding: "5px", borderRadius: "4px" }}
+                >
+                  <DownloadIcon />
+                </Button>
+              </Tooltip>
+            )}
             <Tooltip
               title={
                 sheetRunning ? "Sheet is already running" : "Run the sheet"
@@ -425,6 +449,7 @@ function Sheet(props) {
         hasChanges={hasChanges()}
         onSave={saveSheet}
         sheetRunning={sheetRunning}
+        runId={runId}
       />
       <Box>
         <DataEditor

llmstack/sheets/apis.py

@@ -1,6 +1,9 @@
+import csv
+import io
 import logging
 import uuid
 
+from django.http import StreamingHttpResponse
 from rest_framework import status, viewsets
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response as DRFResponse
@@ -18,6 +21,32 @@
 logger = logging.getLogger(__name__)
 
 
+def _get_sheet_csv(columns, cells, total_rows):
+    output = io.StringIO()
+    writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL)
+
+    # Write header
+    writer.writerow([col.title for col in columns])
+    yield output.getvalue()
+    output.seek(0)
+    output.truncate(0)
+
+    for row in range(1, total_rows + 1):
+        # Create a dict of cell_id to cell from this row
+        row_cells = {cell.cell_id: cell for cell in cells.values() if cell.row == row}
+
+        # Create a list of cell values for this row, using an empty string if the cell doesn't exist
+        row_values = []
+        for column in columns:
+            cell = row_cells.get(f"{column.col}{row}")
+            row_values.append(cell.display_data if cell else "")
+
+        writer.writerow(row_values)
+        yield output.getvalue()
+        output.seek(0)
+        output.truncate(0)
+
+
 class PromptlySheetAppExecuteJob(ProcessingJob):
     @classmethod
     def generate_job_id(cls):
@@ -141,3 +170,39 @@ def run_async(self, request, sheet_uuid=None):
         ).add_to_queue()
 
         return DRFResponse({"job_id": job.id, "run_id": run_entry.uuid}, status=202)
+
+    def download(self, request, sheet_uuid=None):
+        if not sheet_uuid:
+            return DRFResponse(status=status.HTTP_400_BAD_REQUEST)
+
+        profile = Profile.objects.get(user=request.user)
+
+        sheet = PromptlySheet.objects.get(uuid=sheet_uuid, profile_uuid=profile.uuid)
+        if not sheet:
+            return DRFResponse(status=status.HTTP_404_NOT_FOUND)
+
+        response = StreamingHttpResponse(
+            streaming_content=_get_sheet_csv(sheet.columns, sheet.cells, sheet.data.get("total_rows", 0)),
+            content_type="text/csv",
+        )
+        response["Content-Disposition"] = f'attachment; filename="sheet_{sheet_uuid}.csv"'
+        return response
+
+    def download_run(self, request, sheet_uuid=None, run_id=None):
+        if not sheet_uuid or not run_id:
+            return DRFResponse(status=status.HTTP_400_BAD_REQUEST)
+
+        profile = Profile.objects.get(user=request.user)
+
+        sheet = PromptlySheet.objects.get(uuid=sheet_uuid, profile_uuid=profile.uuid)
+        run_entry = PromptlySheetRunEntry.objects.get(uuid=run_id, sheet_uuid=sheet.uuid)
+
+        if not run_entry:
+            return DRFResponse(status=status.HTTP_404_NOT_FOUND)
+
+        response = StreamingHttpResponse(
+            streaming_content=_get_sheet_csv(sheet.columns, sheet.cells, sheet.data.get("total_rows", 0)),
+            content_type="text/csv",
+        )
+        response["Content-Disposition"] = f'attachment; filename="sheet_{sheet_uuid}_{run_id}.csv"'
+        return response

llmstack/sheets/urls.py

@@ -12,4 +12,12 @@
         "api/sheets/<str:sheet_uuid>/run",
         PromptlySheetViewSet.as_view({"post": "run_async"}),
     ),
+    path(
+        "api/sheets/<str:sheet_uuid>/download",
+        PromptlySheetViewSet.as_view({"get": "download"}),
+    ),
+    path(
+        "api/sheets/<str:sheet_uuid>/<str:run_id>/download",
+        PromptlySheetViewSet.as_view({"get": "download_run"}),
+    ),
 ]