BugFix: Step Creation and removal.

* disallowed removing when there is only step
* reliable new step id creation

src/components/model_pipeline/model_pipeline.py

@@ -1,7 +1,6 @@
-import json
-
 import gradio as gr
 import yaml
+from loguru import logger
 
 from components.model_pipeline.state_manager import (
     ModelStepUIState,
@@ -89,6 +88,7 @@ class PipelineInterface:
         step_ui_state: ModelStepUIState,
         available_variables: list[str],
         position: int = 0,
+        n_steps: int = 1,
     ):
         with gr.Column(elem_classes="step-container"):
             # Create the step component
@@ -115,11 +115,14 @@ class PipelineInterface:
             if self.simple:
                 return step_interface
 
+            is_multi_step = n_steps > 1
+            logger.debug(f"Rendering step {position} of {n_steps}")
+
             # Add step controls below
-            with gr.Row(elem_classes="step-controls"):
-                up_button = gr.Button("⬆️ Move Up", elem_classes="step-control-btn")
-                down_button = gr.Button("⬇️ Move Down", elem_classes="step-control-btn")
-                remove_button = gr.Button("🗑️ Remove", elem_classes="step-control-btn")
+            with gr.Row(elem_classes="step-controls", visible=is_multi_step):
+                up_button = gr.Button("⬆️ Move Up", elem_classes="step-control-btn", interactive=is_multi_step)
+                down_button = gr.Button("⬇️ Move Down", elem_classes="step-control-btn", interactive=is_multi_step)
+                remove_button = gr.Button("🗑️ Remove", elem_classes="step-control-btn", interactive=is_multi_step)
 
             buttons = (up_button, down_button, remove_button)
             self._assign_step_controls(buttons, position)
@@ -232,10 +235,10 @@ class PipelineInterface:
 
         # Function to render all steps
         @gr.render(inputs=[self.pipeline_state, self.ui_state])
-        def render_steps(state, ui_state):
+        def render_steps(state: PipelineState, ui_state: PipelineUIState):
             """Render all steps in the pipeline"""
+            logger.info(f"\nRerender triggered! Current UI State:{ui_state.model_dump()}")
             workflow = state.workflow
-            print(f"\nRerender triggered! Current UI State:{ui_state}")
             components = []
 
             step_objects = []  # Reset step objects list
@@ -243,7 +246,7 @@ class PipelineInterface:
                 step_data = workflow.steps[step_id]
                 step_ui_state = ui_state.steps[step_id]
                 available_variables = self.sm.get_all_variables(state, step_id)
-                sub_components = self._render_step(step_data, step_ui_state, available_variables, i)
+                sub_components = self._render_step(step_data, step_ui_state, available_variables, i, ui_state.n_steps)
                 step_objects.append(sub_components)
 
             components.append(step_objects)

src/components/model_pipeline/state_manager.py

@@ -3,6 +3,7 @@ from typing import Any, Literal
 
 import gradio as gr
 import yaml
+from loguru import logger
 from pydantic import BaseModel, Field
 
 from components import utils
@@ -10,17 +11,25 @@ from workflows.factory import create_new_llm_step
 from workflows.structs import ModelStep, Workflow
 
 
-def make_step_id(step_id: int):
+def make_step_id(step_number: int):
     """Make a step id from a step name."""
-    if step_id < 26:
-        return chr(ord("A") + step_id)
+    if step_number < 26:
+        return chr(ord("A") + step_number)
     else:
         # For more than 26 steps, use AA, AB, AC, etc.
-        first_char = chr(ord("A") + (step_id // 26) - 1)
-        second_char = chr(ord("A") + (step_id % 26))
+        first_char = chr(ord("A") + (step_number // 26) - 1)
+        second_char = chr(ord("A") + (step_number % 26))
         return f"{first_char}{second_char}"
 
 
+def make_step_number(step_id: str):
+    """Make a step number from a step id."""
+    if len(step_id) == 1:
+        return ord(step_id) - ord("A")
+    else:
+        return (ord(step_id[0]) - ord("A")) * 26 + (ord(step_id[1]) - ord("A")) + 1
+
+
 class ModelStepUIState(BaseModel):
     """Represents the UI state for a model step component."""
 
@@ -48,6 +57,11 @@ class PipelineUIState(BaseModel):
         """Get the position of a step in the pipeline."""
         return next((i for i, step in enumerate(self.step_ids) if step == step_id), None)
 
+    @property
+    def n_steps(self) -> int:
+        """Get the number of steps in the pipeline."""
+        return len(self.step_ids)
+
     @classmethod
     def from_workflow(cls, workflow: Workflow):
         """Create a pipeline UI state from a workflow."""
@@ -63,7 +77,7 @@ class PipelineState(BaseModel):
     workflow: Workflow
     ui_state: PipelineUIState
 
-    def insert_step(self, position: int, step: ModelStep):
+    def insert_step(self, position: int, step: ModelStep) -> "PipelineState":
         if step.id in self.workflow.steps:
             raise ValueError(f"Step {step.id} already exists in pipeline")
 
@@ -81,14 +95,15 @@ class PipelineState(BaseModel):
             self.ui_state.step_ids.insert(position, step.id)
         return self
 
-    def remove_step(self, position: int):
+    def remove_step(self, position: int) -> "PipelineState":
         step_id = self.ui_state.step_ids.pop(position)
         self.workflow.steps.pop(step_id)
         self.ui_state = self.ui_state.model_copy()
         self.ui_state.steps.pop(step_id)
         self.update_output_variables_mapping()
+        return self
 
-    def update_output_variables_mapping(self):
+    def update_output_variables_mapping(self) -> "PipelineState":
         available_variables = set(self.available_variables)
         for output_field in self.workflow.outputs:
             if self.workflow.outputs[output_field] not in available_variables:
@@ -96,13 +111,21 @@ class PipelineState(BaseModel):
         return self
 
     @property
-    def available_variables(self):
+    def available_variables(self) -> list[str]:
         return self.workflow.get_available_variables()
 
     @property
-    def n_steps(self):
+    def n_steps(self) -> int:
         return len(self.workflow.steps)
 
+    def get_new_step_id(self) -> str:
+        """Get a step ID for a new step."""
+        if not self.workflow.steps:
+            return "A"
+        else:
+            last_step_number = max(map(make_step_number, self.workflow.steps.keys()))
+            return make_step_id(last_step_number + 1)
+
 
 class PipelineStateManager:
     """Manages a pipeline of multiple steps."""
@@ -120,7 +143,7 @@ class PipelineStateManager:
 
     def add_step(self, state: PipelineState, position: int = -1, name=""):
         """Create a new step and return its state."""
-        step_id = make_step_id(state.n_steps)
+        step_id = state.get_new_step_id()
         step_name = name or f"Step {state.n_steps + 1}"
         new_step = create_new_llm_step(step_id=step_id, name=step_name)
         state = state.insert_step(position, new_step)

src/workflows/structs.py

@@ -193,7 +193,8 @@ class Workflow(BaseModel):
 
     def model_dump(self, *args, **kwargs):
         data = super().model_dump(*args, **kwargs)
-        data["steps"] = list(data["steps"].values())
+        if "steps" in data:
+            data["steps"] = list(data["steps"].values())
         return data
 
     @model_validator(mode="before")