Updated datalayers

Author: Allaoua Benchikh

backend/chainlit/data/chainlit_data_layer.py

@@ -338,9 +338,9 @@ async def create_step(self, step_dict: StepDict):
         query = """
         INSERT INTO "Step" (
             id, "threadId", "parentId", input, metadata, name, output,
-            type, "startTime", "endTime", "showInput", "isError"
+            type, "startTime", "endTime", "showInput", "isError", icon
         ) VALUES (
-            $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12
+            $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13
         )
         ON CONFLICT (id) DO UPDATE SET
             "parentId" = COALESCE(EXCLUDED."parentId", "Step"."parentId"),
@@ -359,7 +359,8 @@ async def create_step(self, step_dict: StepDict):
             "endTime" = COALESCE(EXCLUDED."endTime", "Step"."endTime"),
             "startTime" = LEAST(EXCLUDED."startTime", "Step"."startTime"),
             "showInput" = COALESCE(EXCLUDED."showInput", "Step"."showInput"),
-            "isError" = COALESCE(EXCLUDED."isError", "Step"."isError")
+            "isError" = COALESCE(EXCLUDED."isError", "Step"."isError"),
+            icon = COALESCE(EXCLUDED.icon, "Step".icon)
         """
 
         timestamp = await self.get_current_timestamp()
@@ -380,6 +381,7 @@ async def create_step(self, step_dict: StepDict):
             "end_time": timestamp,
             "show_input": str(step_dict.get("showInput", "json")),
             "is_error": step_dict.get("isError", False),
+            "icon": step_dict.get("icon"),
         }
         await self.execute_query(query, params)
 

backend/chainlit/data/literalai.py

@@ -367,6 +367,7 @@ async def create_step(self, step_dict: "StepDict"):
             waitForAnswer=step_dict.get("waitForAnswer"),
             language=step_dict.get("language"),
             showInput=step_dict.get("showInput"),
+            icon=step_dict.get("icon"),
         )
 
         step: LiteralStepDict = {

backend/chainlit/data/sql_alchemy.py

@@ -376,6 +376,8 @@ async def create_step(self, step_dict: "StepDict"):
             if "showInput" in step_dict
             else None
         )
+        # Note: This automatically includes all fields from StepDict including 'icon'
+        # The dynamic column/value generation below will handle any new fields
         parameters = {
             key: value
             for key, value in step_dict.items()

backend/tests/data/test_literalai.py

@@ -83,6 +83,7 @@ def test_step_dict(test_thread) -> StepDict:
         "waitForAnswer": True,
         "showInput": True,
         "language": "en",
+        "icon": "search",
     }
 
 
@@ -179,6 +180,7 @@ async def test_create_step(
                     "waitForAnswer": True,
                     "language": "en",
                     "showInput": True,
+                    "icon": "search",
                 },
                 "input": {"content": "test input"},
                 "output": {"content": "test output"},
@@ -768,6 +770,7 @@ async def test_update_step(
                     "waitForAnswer": True,
                     "language": "en",
                     "showInput": True,
+                    "icon": "search",
                 },
                 "input": {"content": "test input"},
                 "output": {"content": "test output"},

backend/tests/data/test_sql_alchemy.py

@@ -75,6 +75,7 @@ async def data_layer(mock_storage_client: BaseStorageClient, tmp_path: Path):
                     "generation" JSONB,
                     "showInput" TEXT,
                     "language" TEXT,
+                    "icon" TEXT,
                     "indent" INT
                 );
         """