Update sheets api to work with fe

Author: ajhai

llmstack/sheets/apis.py

@@ -1,8 +1,6 @@
 import logging
 import uuid
 
-from asgiref.sync import async_to_sync
-from channels.layers import get_channel_layer
 from rest_framework import status, viewsets
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response as DRFResponse
@@ -14,51 +12,17 @@
     PromptlySheetCell,
     PromptlySheetRunEntry,
 )
-from llmstack.sheets.serializers import PromptlySheetSeializer
-from llmstack.sheets.utils import parse_formula
+from llmstack.sheets.serializers import PromptlySheetSerializer
 
 logger = logging.getLogger(__name__)
 
 
-def write_to_ws_channel(channel_name, message):
-    channel_layer = get_channel_layer()
-    async_to_sync(channel_layer.send)(channel_name, {"type": "send.message", "message": message})
-
-
 class PromptlySheetAppExecuteJob(ProcessingJob):
     @classmethod
     def generate_job_id(cls):
         return "{}".format(str(uuid.uuid4()))
 
 
-def _execute_cell(cell: PromptlySheetCell, sheet: PromptlySheet) -> PromptlySheetCell:
-    if not cell.is_formula:
-        return cell
-
-    # Else execute the formula
-    function_name, params = parse_formula(cell.value)
-    if not function_name or not params:
-        display_value = f"Invalid formula: {cell.value}"
-    else:
-        resolved_params = [""] * len(params)
-        for param_index in range(len(params)):
-            if "-" in params[param_index]:
-                column, row = params[param_index].split("-")
-                resolved_cell = sheet.get_cell(int(row), int(column))
-                resolved_params[param_index] = resolved_cell.value
-
-        display_value = f"Executing {function_name}({', '.join(resolved_params)})"
-
-    extra_data = {**cell.extra_data}
-    extra_data["display_value"] = display_value
-    new_cell = cell.model_copy(update={"extra_data": extra_data})
-    ws_channel_name = sheet.extra_data.get("channel_name")
-    if ws_channel_name:
-        write_to_ws_channel(ws_channel_name, new_cell.model_dump())
-
-    return new_cell
-
-
 class PromptlySheetViewSet(viewsets.ViewSet):
     permission_classes = [IsAuthenticated]
 
@@ -69,7 +33,7 @@ def list(self, request, sheet_uuid=None):
         if sheet_uuid:
             sheet = PromptlySheet.objects.get(uuid=sheet_uuid, profile_uuid=profile.uuid)
             return DRFResponse(
-                PromptlySheetSeializer(
+                PromptlySheetSerializer(
                     instance=sheet,
                     context={
                         "include_cells": include_cells,
@@ -78,7 +42,7 @@ def list(self, request, sheet_uuid=None):
             )
         sheets = PromptlySheet.objects.filter(profile_uuid=profile.uuid)
         return DRFResponse(
-            PromptlySheetSeializer(
+            PromptlySheetSerializer(
                 instance=sheets,
                 many=True,
                 context={
@@ -93,6 +57,7 @@ def create(self, request):
         sheet = PromptlySheet.objects.create(
             name=request.data.get("name"),
             profile_uuid=profile.uuid,
+            data=request.data.get("data", {}),
             extra_data=request.data.get("extra_data", {"has_header": True}),
         )
 
@@ -114,7 +79,7 @@ def create(self, request):
 
             sheet.save(cells=cells_data)
 
-        return DRFResponse(PromptlySheetSeializer(instance=sheet).data)
+        return DRFResponse(PromptlySheetSerializer(instance=sheet).data)
 
     def delete(self, request, sheet_uuid=None):
         profile = Profile.objects.get(user=request.user)
@@ -130,58 +95,50 @@ def patch(self, request, sheet_uuid=None):
         sheet = PromptlySheet.objects.get(uuid=sheet_uuid, profile_uuid=profile.uuid)
         sheet.name = request.data.get("name", sheet.name)
         sheet.extra_data = request.data.get("extra_data", sheet.extra_data)
+
+        if "data" in request.data:
+            if "columns" in request.data["data"]:
+                sheet.data["columns"] = request.data["data"]["columns"]
+
+            if "total_rows" in request.data["data"]:
+                sheet.data["total_rows"] = request.data["data"]["total_rows"]
+
+            if "description" in request.data["data"]:
+                sheet.data["description"] = request.data["data"]["description"]
+
         sheet.save()
 
-        if "cells" in request.data:
+        if "data" in request.data and "cells" in request.data["data"]:
             cells_data = []
-            for row in range(len(request.data["cells"])):
-                cells_row_data = []
-                for column in range(len(request.data["cells"][row])):
-                    cell_data = request.data["cells"][row][column]
-                    cells_row_data.append(PromptlySheetCell(row=row, col=column, **cell_data))
-                cells_data.append(cells_row_data)
+            for row_number in request.data["data"]["cells"]:
+                if not isinstance(request.data["data"]["cells"][row_number], dict):
+                    raise ValueError("Invalid cells data")
+                for column_number in request.data["data"]["cells"][row_number]:
+                    cell_data = request.data["data"]["cells"][row_number][column_number]
+                    cell_data.pop("row", None)
+                    cell_data.pop("col", None)
+                    cells_data.append(PromptlySheetCell(row=row_number, col=column_number, **cell_data))
             sheet.save(cells=cells_data)
 
-        return DRFResponse(PromptlySheetSeializer(instance=sheet).data)
+        return DRFResponse(PromptlySheetSerializer(instance=sheet).data)
 
-    def execute(self, request, sheet_uuid=None):
+    def run_async(self, request, sheet_uuid=None):
         profile = Profile.objects.get(user=request.user)
-
         sheet = PromptlySheet.objects.get(uuid=sheet_uuid, profile_uuid=profile.uuid)
         if sheet.is_locked:
             return DRFResponse(
-                {"detail": "The sheet is locked and cannot be executed"},
+                {"detail": "The sheet is locked and cannot be run at this time."},
                 status=status.HTTP_400_BAD_REQUEST,
             )
 
         sheet.is_locked = True
         sheet.save(update_fields=["is_locked"])
 
-        try:
-            processed_cells = {}
-            for row_number, column_cells in sheet.rows:
-                processed_cells_row = {}
-                for column_number in column_cells:
-                    processed_cells_row[column_number] = _execute_cell(column_cells[column_number], sheet)
-                processed_cells[row_number] = processed_cells_row
-
-            if processed_cells:
-                sheet.save(cells=processed_cells, update_fields=["updated_at"])
-                # Store the processed data in sheet runs table
-                run_entry = PromptlySheetRunEntry(sheet_uuid=sheet.uuid, profile_uuid=profile.uuid)
-                run_entry.save(cells=processed_cells)
-
-        except Exception:
-            logger.exception("Error executing sheet")
-
-        sheet.is_locked = False
-        sheet.save(update_fields=["is_locked"])
-        return DRFResponse(PromptlySheetSeializer(instance=sheet).data)
+        run_entry = PromptlySheetRunEntry(sheet_uuid=sheet.uuid, profile_uuid=profile.uuid)
 
-    def execute_async(self, request, sheet_uuid=None):
         job = PromptlySheetAppExecuteJob.create(
-            func="llmstack.promptly_sheets.tasks.process_sheet_execute_request",
-            args=[request.user.email, sheet_uuid],
+            func="llmstack.sheets.tasks.run_sheet",
+            args=[sheet, run_entry],
         ).add_to_queue()
 
-        return DRFResponse({"job_id": job.id}, status=202)
+        return DRFResponse({"job_id": job.id, "run_id": run_entry.uuid}, status=202)

llmstack/sheets/models.py

@@ -1,7 +1,8 @@
 import base64
 import json
+import logging
 import uuid
-from typing import Dict
+from typing import List
 
 from django.db import models
 from django.db.models.signals import post_delete
@@ -10,17 +11,19 @@
 
 from llmstack.assets.models import Assets
 
+logger = logging.getLogger(__name__)
+
 
 class PromptlySheetCell(BaseModel):
     row: int
     col: int
-    value: str = ""
-    value_type: str = "string"
-    extra_data: dict = {}
+    data: str = ""
+    formula: str = ""
+    kind: str = "string"
 
     @property
     def is_formula(self):
-        return self.value.startswith("=")
+        return bool(self.formula)
 
     @property
     def cell_id(self):
@@ -94,13 +97,17 @@ def delete_sheet_data_objrefs(data_objrefs):
             pass
 
 
-def create_sheet_data_objrefs(
-    cells: Dict[int, Dict[int, PromptlySheetCell]], sheet_name, sheet_uuid, page_size: int = 1000
-):
-    for i in range(0, len(cells), page_size):
+def create_sheet_data_objrefs(cells: List[PromptlySheetCell], sheet_name, sheet_uuid, page_size: int = 1000):
+    # Sort the cells by row and columns
+    cells = sorted(cells, key=lambda cell: (cell.row, cell.col))
+    max_rows = max(cells, key=lambda cell: cell.row).row + 1
+
+    for i in range(0, max_rows, page_size):
         chunk = {}
-        for j in range(i, min(i + page_size, len(cells))):
-            chunk[j] = dict(map(lambda entry: (entry[0], entry[1].model_dump()), cells[j].items()))
+        for j in range(i, min(i + page_size, max_rows)):
+            # Find cells from this row and add them to the chunk
+            cells_in_row = list(filter(lambda cell: cell.row == j, cells))
+            chunk[j] = dict(map(lambda cell: (cell.col, cell.model_dump()), cells_in_row))
 
         data_json = json.dumps(chunk)
         filename = f"{sheet_name}_{str(uuid.uuid4())[:4]}_{i}.json"
@@ -135,9 +142,12 @@ def get_sheet_cells(objref):
                 )
             )
 
-    except Exception:
+    except Exception as e:
+        logger.error(f"Error loading sheet data from objref: {e}")
         pass
 
+    return {}
+
 
 class PromptlySheet(models.Model):
     uuid = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
@@ -178,19 +188,22 @@ def rows(self):
             for cell_row in cells:
                 yield (cell_row, cells[cell_row])
 
+    @property
+    def columns(self):
+        return self.data.get("columns", [])
+
     def save(self, *args, **kwargs):
-        if kwargs.get("cells"):
+        if "cells" in kwargs:
             if self.data and self.data.get("cells"):
                 # Delete the older objrefs
                 delete_sheet_data_objrefs(self.data.get("cells", []))
             cell_objs = kwargs.pop("cells")
-            self.data = {
-                "cells": list(
-                    create_sheet_data_objrefs(
-                        cell_objs, self.name, str(self.uuid), page_size=self.extra_data.get("page_size", 1000)
-                    )
+
+            self.data["cells"] = list(
+                create_sheet_data_objrefs(
+                    cell_objs, self.name, str(self.uuid), page_size=self.extra_data.get("page_size", 1000)
                 )
-            }
+            )
             if kwargs.get("update_fields"):
                 kwargs["update_fields"].append("data")
 

llmstack/sheets/serializers.py

@@ -1,10 +1,17 @@
+import logging
+
 from rest_framework import serializers
 
 from llmstack.sheets.models import PromptlySheet
 
+logger = logging.getLogger(__name__)
+
 
-class PromptlySheetSeializer(serializers.ModelSerializer):
+class PromptlySheetSerializer(serializers.ModelSerializer):
     cells = serializers.SerializerMethodField()
+    columns = serializers.SerializerMethodField()
+    description = serializers.SerializerMethodField()
+    total_rows = serializers.SerializerMethodField()
 
     def get_cells(self, obj):
         cells = {}
@@ -15,6 +22,25 @@ def get_cells(self, obj):
 
         return cells
 
+    def get_columns(self, obj):
+        return obj.data.get("columns", [])
+
+    def get_description(self, obj):
+        return obj.data.get("description", "")
+
+    def get_total_rows(self, obj):
+        return obj.data.get("total_rows", 0)
+
     class Meta:
         model = PromptlySheet
-        fields = ["uuid", "name", "extra_data", "cells", "created_at", "updated_at"]
+        fields = [
+            "uuid",
+            "name",
+            "extra_data",
+            "cells",
+            "columns",
+            "total_rows",
+            "description",
+            "created_at",
+            "updated_at",
+        ]

llmstack/sheets/urls.py

@@ -9,11 +9,7 @@
         PromptlySheetViewSet.as_view({"get": "list", "patch": "patch", "delete": "delete"}),
     ),
     path(
-        "api/sheets/<str:sheet_uuid>/execute",
-        PromptlySheetViewSet.as_view({"post": "execute"}),
-    ),
-    path(
-        "api/sheets/<str:sheet_uuid>/execute_async",
-        PromptlySheetViewSet.as_view({"post": "execute_async"}),
+        "api/sheets/<str:sheet_uuid>/run",
+        PromptlySheetViewSet.as_view({"post": "run_async"}),
     ),
 ]