Execute app_runs

ajhai · ajhai · commit dc4c0476f695 · 2024-08-24T00:06:21.000-07:00
diff --git a/llmstack/sheets/apis.py b/llmstack/sheets/apis.py
@@ -10,6 +10,7 @@
 from llmstack.sheets.models import (
     PromptlySheet,
     PromptlySheetCell,
+    PromptlySheetColumn,
     PromptlySheetRunEntry,
 )
 from llmstack.sheets.serializers import PromptlySheetSerializer
@@ -98,7 +99,9 @@ def patch(self, request, sheet_uuid=None):
 
         if "data" in request.data:
             if "columns" in request.data["data"]:
-                sheet.data["columns"] = request.data["data"]["columns"]
+                sheet.data["columns"] = [
+                    PromptlySheetColumn(**column_data).model_dump() for column_data in request.data["data"]["columns"]
+                ]
 
             if "total_rows" in request.data["data"]:
                 sheet.data["total_rows"] = request.data["data"]["total_rows"]
@@ -138,7 +141,7 @@ def run_async(self, request, sheet_uuid=None):
 
         job = PromptlySheetAppExecuteJob.create(
             func="llmstack.sheets.tasks.run_sheet",
-            args=[sheet, run_entry],
+            args=[sheet, run_entry, request.user],
         ).add_to_queue()
 
         return DRFResponse({"job_id": job.id, "run_id": run_entry.uuid}, status=202)
diff --git a/llmstack/sheets/consumers.py b/llmstack/sheets/consumers.py
@@ -4,13 +4,13 @@
 from channels.db import database_sync_to_async
 from channels.generic.websocket import AsyncWebsocketConsumer
 
-from llmstack.sheets.models import PromptlySheet
-
 logger = logging.getLogger(__name__)
 
 
 @database_sync_to_async
 def update_sheet_channel_name(sheet_id, channel_name):
+    from llmstack.sheets.models import PromptlySheet
+
     # Update the channel name for the sheet
     sheet = PromptlySheet.objects.get(uuid=sheet_id)
     sheet_extra_data = sheet.extra_data or {}
@@ -22,18 +22,21 @@ def update_sheet_channel_name(sheet_id, channel_name):
 class SheetAppConsumer(AsyncWebsocketConsumer):
     async def connect(self):
         self.sheet_id = self.scope["url_route"]["kwargs"]["sheet_id"]
+        self.run_id = self.scope["url_route"]["kwargs"]["run_id"]
         self._user = self.scope.get("user", None)
         # Update the channel name for the sheet
-        await update_sheet_channel_name(self.sheet_id, self.channel_name)
+        await self.channel_layer.group_add(self.run_id, self.channel_name)
         await super().connect()
 
     async def disconnect(self, close_code):
-        await update_sheet_channel_name(self.sheet_id, None)
+        await self.channel_layer.group_discard(self.run_id, self.channel_name)
 
     async def close(self, code=None, reason=None):
         await update_sheet_channel_name(self.sheet_id, None)
         await super().close(code, reason)
 
-    # Called when a message is received on the channel
-    async def send_message(self, event):
+    async def cell_update(self, event):
+        await self.send(text_data=json.dumps(event))
+
+    async def cell_updating(self, event):
         await self.send(text_data=json.dumps(event))
diff --git a/llmstack/sheets/models.py b/llmstack/sheets/models.py
@@ -2,7 +2,8 @@
 import json
 import logging
 import uuid
-from typing import List
+from enum import Enum
+from typing import List, Optional
 
 from django.db import models
 from django.db.models.signals import post_delete
@@ -14,20 +15,51 @@
 logger = logging.getLogger(__name__)
 
 
+class PromptlySheetColumnType(str, Enum):
+    TEXT = "text"
+    NUMBER = "number"
+    BOOLEAN = "boolean"
+    DATE = "date"
+    TIME = "time"
+    DATETIME = "datetime"
+    DURATION = "duration"
+    PERCENTAGE = "percentage"
+    CURRENCY = "currency"
+    URL = "url"
+    EMAIL = "email"
+    PHONE = "phone"
+    IMAGE = "image"
+    FILE = "file"
+    FORMULA = "formula"
+    APP_RUN = "app_run"
+    PROCESSOR_RUN = "processor_run"
+
+    def __str__(self):
+        return self.value
+
+
+class PromptlySheetColumn(BaseModel):
+    title: str
+    kind: PromptlySheetColumnType = PromptlySheetColumnType.TEXT
+    data: dict = {}
+    col: int
+    width: Optional[int] = None
+
+
 class PromptlySheetCell(BaseModel):
     row: int
     col: int
     data: str = ""
     formula: str = ""
-    kind: str = "string"
+    kind: PromptlySheetColumnType = PromptlySheetColumnType.TEXT
 
     @property
     def is_formula(self):
         return bool(self.formula)
 
     @property
     def cell_id(self):
-        return f"{self.col}-{self.row}"
+        return f"{self.row}-{self.col}"
 
     def model_dump(
         self,
@@ -132,7 +164,10 @@ def get_sheet_cells(objref):
                         (
                             dict(
                                 map(
-                                    lambda cell_entry: (int(cell_entry[0]), PromptlySheetCell(**cell_entry[1])),
+                                    lambda cell_entry: (
+                                        int(cell_entry[0]),
+                                        PromptlySheetCell(**cell_entry[1]),
+                                    ),
                                     row_entry[1].items(),
                                 )
                             )
@@ -185,12 +220,12 @@ def cells(self):
     def rows(self):
         for objref in self.data.get("cells", []):
             cells = get_sheet_cells(objref)
-            for cell_row in cells:
-                yield (cell_row, cells[cell_row])
+            for row, row_cells in cells.items():
+                yield (row, row_cells)
 
     @property
     def columns(self):
-        return self.data.get("columns", [])
+        return [PromptlySheetColumn(**column) for column in self.data.get("columns", [])]
 
     def save(self, *args, **kwargs):
         if "cells" in kwargs:
diff --git a/llmstack/sheets/tasks.py b/llmstack/sheets/tasks.py
@@ -1,17 +1,117 @@
-from llmstack.sheets.apis import PromptlySheetViewSet
+import logging
+import uuid
 
+from asgiref.sync import async_to_sync
+from channels.layers import get_channel_layer
+from django.contrib.auth.models import User
+from django.test import RequestFactory
+from rest_framework.response import Response as DRFResponse
 
-def process_sheet_execute_request(user_email, sheet_id):
-    from django.contrib.auth.models import User
-    from django.test import RequestFactory
+from llmstack.common.utils.utils import hydrate_input
+from llmstack.sheets.models import PromptlySheet, PromptlySheetCell, PromptlySheetColumn
+from llmstack.sheets.serializers import PromptlySheetSerializer
 
-    user = User.objects.get(email=user_email)
+try:
+    from promptly.promptly_app_store.apis import AppStoreAppViewSet
+except ImportError:
+    from llmstack.app_store.apis import AppStoreAppViewSet
+
+
+logger = logging.getLogger(__name__)
+channel_layer = get_channel_layer()
+
+
+def number_to_letters(num):
+    letters = ""
+    while num >= 0:
+        letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[num % 26] + letters
+        num = num // 26 - 1
+    return letters
+
+
+def _execute_cell(
+    cell: PromptlySheetCell,
+    column: PromptlySheetColumn,
+    row: dict[str, PromptlySheetCell],
+    sheet: PromptlySheet,
+    run_id: str,
+    user: User,
+) -> PromptlySheetCell:
+    if column.kind != "app_run":
+        return cell
+
+    async_to_sync(channel_layer.group_send)(run_id, {"type": "cell.updating", "cell": {"id": cell.cell_id}})
+
+    app_slug = column.data["app_slug"]
+    input_values = {number_to_letters(col): cell.data for col, cell in row.items()}
+    input = hydrate_input(column.data["input"], input_values)
 
     request = RequestFactory().post(
-        f"/api/sheets/{sheet_id}/execute",
+        f"/api/store/apps/{app_slug}",
         format="json",
     )
+    request.data = {
+        "stream": False,
+        "input": input,
+    }
     request.user = user
-    response = PromptlySheetViewSet().execute(request, sheet_uuid=sheet_id)
 
-    return {"status_code": response.status_code, "data": response}
+    # Execute the app
+    response = async_to_sync(AppStoreAppViewSet().run_app_internal_async)(
+        slug=app_slug,
+        session_id=None,
+        request_uuid=str(uuid.uuid4()),
+        request=request,
+    )
+
+    output = response.get("output", "")
+    async_to_sync(channel_layer.group_send)(
+        run_id, {"type": "cell.update", "cell": {"id": cell.cell_id, "data": response.get("output", "")}}
+    )
+    cell.data = output
+
+    return cell
+
+
+def run_sheet(sheet, run_entry, user):
+    try:
+        processed_cells = []
+        existing_rows = list(sheet.rows)
+        existing_cols = sheet.columns
+
+        for row_number in range(sheet.data.get("total_rows", 0)):
+            # Find existing row and the cell
+            existing_row = next(filter(lambda row: row[0] == row_number, existing_rows), None)
+            existing_row_cells = existing_row[1] if existing_row else {}
+
+            for column in existing_cols:
+                if column.kind != "app_run":
+                    if existing_row:
+                        existing_cell = next(
+                            filter(lambda cell: cell.col == column.col, existing_row_cells.values()), None
+                        )
+                        if existing_cell:
+                            processed_cells.append(existing_cell)
+                    continue
+
+                # Create a new cell
+                cell_to_execute = PromptlySheetCell(
+                    row=row_number,
+                    col=column.col,
+                    kind=column.kind,
+                )
+                processed_cells.append(
+                    _execute_cell(cell_to_execute, column, existing_row_cells, sheet, str(run_entry.uuid), user)
+                )
+
+        if processed_cells:
+            sheet.save(cells=processed_cells, update_fields=["updated_at"])
+            # Store the processed data in sheet runs table
+            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(PromptlySheetSerializer(instance=sheet).data)
