From d55381081ac0632d23932ec4d8c5339c1d448a74 Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <46546486+ryanhopperlowe@users.noreply.github.com>
Date: Tue, 31 Dec 2024 17:09:01 -0600
Subject: [PATCH] enhance: update workflow step UX (#1087)

- adding a step auto-adds a generic step. No longer need to specify step type when adding
- step types can be updated via dropdown on the step list item
---
 ui/admin/app/components/ui/textarea.tsx       |   6 +-
 .../app/components/workflow/steps/AddStep.tsx | 135 +-----------
 ui/admin/app/components/workflow/steps/If.tsx | 154 -------------
 .../components/workflow/steps/IfContent.tsx   |  79 +++++++
 .../app/components/workflow/steps/Step.tsx    | 208 ------------------
 .../components/workflow/steps/StepBase.tsx    | 116 ++++++++++
 .../components/workflow/steps/StepContent.tsx | 127 +++++++++++
 .../workflow/steps/StepRenderer.tsx           |  55 +++--
 .../workflow/steps/StepTypeSelect.tsx         |  65 ++++++
 .../app/components/workflow/steps/While.tsx   | 153 -------------
 .../workflow/steps/WhileContent.tsx           |  83 +++++++
 ui/admin/app/lib/model/workflows.ts           |  40 +++-
 12 files changed, 547 insertions(+), 674 deletions(-)
 delete mode 100644 ui/admin/app/components/workflow/steps/If.tsx
 create mode 100644 ui/admin/app/components/workflow/steps/IfContent.tsx
 delete mode 100644 ui/admin/app/components/workflow/steps/Step.tsx
 create mode 100644 ui/admin/app/components/workflow/steps/StepBase.tsx
 create mode 100644 ui/admin/app/components/workflow/steps/StepContent.tsx
 create mode 100644 ui/admin/app/components/workflow/steps/StepTypeSelect.tsx
 delete mode 100644 ui/admin/app/components/workflow/steps/While.tsx
 create mode 100644 ui/admin/app/components/workflow/steps/WhileContent.tsx

diff --git a/ui/admin/app/components/ui/textarea.tsx b/ui/admin/app/components/ui/textarea.tsx
index b99fb0459..79d76c0fc 100644
--- a/ui/admin/app/components/ui/textarea.tsx
+++ b/ui/admin/app/components/ui/textarea.tsx
@@ -4,12 +4,14 @@ import { forwardRef, useImperativeHandle } from "react";
 
 import { cn } from "~/lib/utils";
 
+// note: use outline instead of ring to avoid overriding the ring from the outline variant
 const textareaVariants = cva(
-    "flex w-full rounded-md bg-transparent text-sm placeholder:text-muted-foreground has-[:focus-visible]:ring-1 has-[:focus-visible]:ring-ring group group-disabled:cursor-not-allowed group-disabled:bg-opacity-50",
+    "flex w-full rounded-md bg-transparent text-sm placeholder:text-muted-foreground has-[:focus-visible]:outline has-[:focus-visible]:outline-1 has-[:focus-visible]:outline-ring group group-disabled:cursor-not-allowed group-disabled:bg-opacity-50",
     {
         variants: {
             variant: {
-                outlined: "border border-input shadow-sm",
+                // note: use inset ring instead of border so that the wrapper doesn't add any extra height or width
+                outlined: "ring-1 ring-inset ring-input",
                 flat: "border-none shadow-none bg-muted",
             },
         },
diff --git a/ui/admin/app/components/workflow/steps/AddStep.tsx b/ui/admin/app/components/workflow/steps/AddStep.tsx
index 1eea3f0f2..4a4bd8325 100644
--- a/ui/admin/app/components/workflow/steps/AddStep.tsx
+++ b/ui/admin/app/components/workflow/steps/AddStep.tsx
@@ -1,30 +1,8 @@
-import { ArrowRight, GitFork, Plus, RotateCw } from "lucide-react";
-import { useMemo, useState } from "react";
-import useSWR from "swr";
+import { PlusIcon } from "lucide-react";
 
-import {
-    ToolReference,
-    toolReferenceToTemplate,
-} from "~/lib/model/toolReferences";
 import { Step, getDefaultStep } from "~/lib/model/workflows";
-import { ToolReferenceService } from "~/lib/service/api/toolreferenceService";
 
 import { Button } from "~/components/ui/button";
-import {
-    Dialog,
-    DialogContent,
-    DialogHeader,
-    DialogTitle,
-} from "~/components/ui/dialog";
-import {
-    Popover,
-    PopoverContent,
-    PopoverTrigger,
-} from "~/components/ui/popover";
-import { ScrollArea } from "~/components/ui/scroll-area";
-import { StepTemplateCard } from "~/components/workflow/steps/StepTemplateCard";
-
-type StepType = "regular" | "if" | "while" | "template";
 
 interface AddStepButtonProps {
     onAddStep: (newStep: Step) => void;
@@ -32,109 +10,14 @@ interface AddStepButtonProps {
 }
 
 export function AddStepButton({ onAddStep, className }: AddStepButtonProps) {
-    const [open, setOpen] = useState(false);
-    const [isTemplateModalOpen, setIsTemplateModalOpen] = useState(false);
-
-    const getStepTemplates = useSWR(
-        ToolReferenceService.getToolReferences.key("stepTemplate"),
-        () => ToolReferenceService.getToolReferences("stepTemplate")
-    );
-
-    const stepTemplates = useMemo(() => {
-        if (!getStepTemplates.data) return [];
-        return getStepTemplates.data;
-    }, [getStepTemplates.data]);
-
-    const createNewStep = (type: StepType) => {
-        if (type === "template") {
-            setIsTemplateModalOpen(true);
-        } else {
-            const newStep = getDefaultStep();
-
-            if (type === "if") {
-                newStep.if = { condition: "", steps: [], else: [] };
-            } else if (type === "while") {
-                newStep.while = { condition: "", maxLoops: 0, steps: [] };
-            }
-
-            onAddStep(newStep);
-            setOpen(false);
-        }
-    };
-
-    const handleStepTemplateSelection = (stepTemplate: ToolReference) => {
-        const newStep: Step = {
-            id: Date.now().toString(),
-            name: stepTemplate.name,
-            description: "",
-            template: toolReferenceToTemplate(stepTemplate),
-            cache: false,
-            temperature: 0,
-            tools: [],
-            agents: [],
-            workflows: [],
-        };
-        onAddStep(newStep);
-        setIsTemplateModalOpen(false);
-    };
-
     return (
-        <>
-            <Popover open={open} onOpenChange={setOpen}>
-                <PopoverTrigger asChild>
-                    <Button variant="ghost" className={className}>
-                        <Plus /> Add Step
-                    </Button>
-                </PopoverTrigger>
-
-                <PopoverContent
-                    className="w-22 bg-secondary dark:bg-zinc-800 shadow-2xl"
-                    side="top"
-                >
-                    <div className="grid gap-4">
-                        <Button
-                            className="dark:bg-zinc-600 dark:text-white"
-                            onClick={() => createNewStep("while")}
-                        >
-                            <RotateCw className="w-4 h-4 mr-2" /> While
-                        </Button>
-                        <Button
-                            className="dark:bg-zinc-600 dark:text-white"
-                            onClick={() => createNewStep("if")}
-                        >
-                            <GitFork className="w-4 h-4 mr-2" /> If
-                        </Button>
-                        <Button
-                            className="dark:bg-zinc-600 dark:text-white"
-                            onClick={() => createNewStep("regular")}
-                        >
-                            <ArrowRight className="w-4 h-4 mr-2" /> Step
-                        </Button>
-                    </div>
-                </PopoverContent>
-            </Popover>
-
-            <Dialog
-                open={isTemplateModalOpen}
-                onOpenChange={setIsTemplateModalOpen}
-            >
-                <DialogContent className="min-w-[50vw] max-h-[80vh]">
-                    <DialogHeader>
-                        <DialogTitle>Select a Template</DialogTitle>
-                    </DialogHeader>
-                    <ScrollArea className="grid gap-4 h-[70vh]">
-                        {stepTemplates.map((template, index) => (
-                            <StepTemplateCard
-                                key={index}
-                                stepTemplate={template}
-                                onClick={() =>
-                                    handleStepTemplateSelection(template)
-                                }
-                            />
-                        ))}
-                    </ScrollArea>
-                </DialogContent>
-            </Dialog>
-        </>
+        <Button
+            variant="ghost"
+            className={className}
+            startContent={<PlusIcon />}
+            onClick={() => onAddStep(getDefaultStep())}
+        >
+            Add Step
+        </Button>
     );
 }
diff --git a/ui/admin/app/components/workflow/steps/If.tsx b/ui/admin/app/components/workflow/steps/If.tsx
deleted file mode 100644
index 8079985bd..000000000
--- a/ui/admin/app/components/workflow/steps/If.tsx
+++ /dev/null
@@ -1,154 +0,0 @@
-import { ChevronDown, ChevronRight, GitFork, Trash } from "lucide-react";
-import { useState } from "react";
-
-import { If, Step } from "~/lib/model/workflows";
-import { cn } from "~/lib/utils";
-
-import { Button } from "~/components/ui/button";
-import { ClickableDiv } from "~/components/ui/clickable-div";
-import { AutosizeTextarea } from "~/components/ui/textarea";
-import { AddStepButton } from "~/components/workflow/steps/AddStep";
-
-export function IfComponent({
-    ifCondition,
-    onUpdate,
-    onDelete,
-    renderStep,
-    className,
-}: {
-    ifCondition: If;
-    onUpdate: (updatedIf: If) => void;
-    onDelete: () => void;
-    renderStep: (
-        step: Step,
-        onUpdate: (updatedStep: Step) => void,
-        onDelete: () => void
-    ) => React.ReactNode;
-    className?: string;
-}) {
-    const [isExpanded, setIsExpanded] = useState(false);
-
-    const addStep = (branch: "steps" | "else") => (newStep: Step) => {
-        onUpdate({
-            ...ifCondition,
-            [branch]: [...(ifCondition[branch] || []), newStep],
-        });
-    };
-
-    const updateNestedStep = (
-        branch: "steps" | "else",
-        index: number,
-        updatedStep: Step
-    ) => {
-        const newBranch = [...(ifCondition[branch] || [])];
-        newBranch[index] = updatedStep;
-        onUpdate({ ...ifCondition, [branch]: newBranch });
-    };
-
-    const deleteNestedStep = (branch: "steps" | "else", index: number) => {
-        const newBranch = (ifCondition[branch] || []).filter(
-            (_, i) => i !== index
-        );
-        onUpdate({ ...ifCondition, [branch]: newBranch });
-    };
-
-    return (
-        <div className={cn("border rounded-md", className)}>
-            <ClickableDiv
-                className={cn(
-                    "flex items-start p-3 bg-secondary",
-                    isExpanded ? "rounded-t-md" : "rounded-md"
-                )}
-                onClick={() => setIsExpanded((prev) => !prev)}
-            >
-                <div className="flex items-center">
-                    <Button
-                        variant="ghost"
-                        size="icon"
-                        className="p-0 w-6 h-6 mr-2 self-center"
-                    >
-                        {isExpanded ? (
-                            <ChevronDown className="w-4 h-4" />
-                        ) : (
-                            <ChevronRight className="w-4 h-4" />
-                        )}
-                    </Button>
-
-                    <div className="flex items-center justify-center p-2 w-24 border bg-background rounded-md mr-2">
-                        <GitFork className="w-4 h-4 mr-1" />
-                        <span className="text-sm font-medium">If</span>
-                    </div>
-                </div>
-
-                <AutosizeTextarea
-                    value={ifCondition.condition || ""}
-                    onChange={(e) =>
-                        onUpdate({ ...ifCondition, condition: e.target.value })
-                    }
-                    onClick={(e) => e.stopPropagation()}
-                    minHeight={0}
-                    maxHeight={100}
-                    placeholder="Condition"
-                    className="flex-grow bg-background"
-                />
-
-                <Button
-                    variant="ghost"
-                    size="icon"
-                    onClick={(e) => {
-                        e.stopPropagation();
-                        onDelete();
-                    }}
-                    className="ml-2"
-                >
-                    <Trash className="w-4 h-4" />
-                </Button>
-            </ClickableDiv>
-
-            {isExpanded && (
-                <div className="p-3 space-y-4">
-                    <div className="space-y-2">
-                        <h4 className="font-semibold">Then:</h4>
-                        {ifCondition.steps?.map((step, index) => (
-                            <div key={index} className="ml-4">
-                                {renderStep(
-                                    step,
-                                    (updatedStep) =>
-                                        updateNestedStep(
-                                            "steps",
-                                            index,
-                                            updatedStep
-                                        ),
-                                    () => deleteNestedStep("steps", index)
-                                )}
-                            </div>
-                        ))}
-                        <div className="ml-4">
-                            <AddStepButton onAddStep={addStep("steps")} />
-                        </div>
-                    </div>
-                    <div className="space-y-2">
-                        <h4 className="font-semibold">Else:</h4>
-                        {ifCondition.else?.map((step, index) => (
-                            <div key={index} className="ml-4">
-                                {renderStep(
-                                    step,
-                                    (updatedStep) =>
-                                        updateNestedStep(
-                                            "else",
-                                            index,
-                                            updatedStep
-                                        ),
-                                    () => deleteNestedStep("else", index)
-                                )}
-                            </div>
-                        ))}
-                        <div className="ml-4">
-                            <AddStepButton onAddStep={addStep("else")} />
-                        </div>
-                    </div>
-                </div>
-            )}
-        </div>
-    );
-}
diff --git a/ui/admin/app/components/workflow/steps/IfContent.tsx b/ui/admin/app/components/workflow/steps/IfContent.tsx
new file mode 100644
index 000000000..fe50a2cc0
--- /dev/null
+++ b/ui/admin/app/components/workflow/steps/IfContent.tsx
@@ -0,0 +1,79 @@
+import { If, Step } from "~/lib/model/workflows";
+
+import { AddStepButton } from "~/components/workflow/steps/AddStep";
+
+export function IfContent({
+    ifCondition,
+    onUpdate,
+    renderStep,
+}: {
+    ifCondition: If;
+    onUpdate: (updatedIf: If) => void;
+    renderStep: (
+        step: Step,
+        onUpdate: (updatedStep: Step) => void,
+        onDelete: () => void
+    ) => React.ReactNode;
+    className?: string;
+}) {
+    const addStep = (branch: "steps" | "else") => (newStep: Step) => {
+        onUpdate({
+            ...ifCondition,
+            [branch]: [...(ifCondition[branch] || []), newStep],
+        });
+    };
+
+    const updateNestedStep = (
+        branch: "steps" | "else",
+        index: number,
+        updatedStep: Step
+    ) => {
+        const newBranch = [...(ifCondition[branch] || [])];
+        newBranch[index] = updatedStep;
+        onUpdate({ ...ifCondition, [branch]: newBranch });
+    };
+
+    const deleteNestedStep = (branch: "steps" | "else", index: number) => {
+        const newBranch = (ifCondition[branch] || []).filter(
+            (_, i) => i !== index
+        );
+        onUpdate({ ...ifCondition, [branch]: newBranch });
+    };
+
+    return (
+        <div className="p-3 space-y-4">
+            <div className="space-y-2">
+                <h4 className="font-semibold">Then:</h4>
+                {ifCondition.steps?.map((step, index) => (
+                    <div key={index} className="ml-4">
+                        {renderStep(
+                            step,
+                            (updatedStep) =>
+                                updateNestedStep("steps", index, updatedStep),
+                            () => deleteNestedStep("steps", index)
+                        )}
+                    </div>
+                ))}
+                <div className="ml-4">
+                    <AddStepButton onAddStep={addStep("steps")} />
+                </div>
+            </div>
+            <div className="space-y-2">
+                <h4 className="font-semibold">Else:</h4>
+                {ifCondition.else?.map((step, index) => (
+                    <div key={index} className="ml-4">
+                        {renderStep(
+                            step,
+                            (updatedStep) =>
+                                updateNestedStep("else", index, updatedStep),
+                            () => deleteNestedStep("else", index)
+                        )}
+                    </div>
+                ))}
+                <div className="ml-4">
+                    <AddStepButton onAddStep={addStep("else")} />
+                </div>
+            </div>
+        </div>
+    );
+}
diff --git a/ui/admin/app/components/workflow/steps/Step.tsx b/ui/admin/app/components/workflow/steps/Step.tsx
deleted file mode 100644
index 80d487aa8..000000000
--- a/ui/admin/app/components/workflow/steps/Step.tsx
+++ /dev/null
@@ -1,208 +0,0 @@
-import {
-    ArrowRight,
-    ChevronDown,
-    ChevronRight,
-    CogIcon,
-    Puzzle,
-    Trash,
-    User,
-    Wrench,
-} from "lucide-react";
-import { useState } from "react";
-
-import { Step } from "~/lib/model/workflows";
-import { cn } from "~/lib/utils";
-
-import { AgentSelectModule } from "~/components/agent/shared/AgentSelect";
-import { BasicToolForm } from "~/components/tools/BasicToolForm";
-import {
-    Accordion,
-    AccordionContent,
-    AccordionItem,
-    AccordionTrigger,
-} from "~/components/ui/accordion";
-import { Button } from "~/components/ui/button";
-import { ClickableDiv } from "~/components/ui/clickable-div";
-import { Input } from "~/components/ui/input";
-import { Switch } from "~/components/ui/switch";
-import { AutosizeTextarea } from "~/components/ui/textarea";
-import { WorkflowSelectModule } from "~/components/workflow/WorkflowSelectModule";
-
-export function StepComponent({
-    step,
-    onUpdate,
-    onDelete,
-    className,
-}: {
-    step: Step;
-    onUpdate: (updatedStep: Step) => void;
-    onDelete: () => void;
-    className?: string;
-}) {
-    const [isExpanded, setIsExpanded] = useState(false);
-
-    return (
-        <div className={cn("border rounded-md", className)}>
-            <ClickableDiv
-                className={cn(
-                    "flex items-start p-3 bg-secondary",
-                    isExpanded ? "rounded-t-md" : "rounded-md"
-                )}
-                onClick={() => setIsExpanded((prev) => !prev)}
-            >
-                <div className="flex items-center">
-                    <Button
-                        variant="ghost"
-                        size="icon"
-                        className="p-0 w-6 h-6 mr-2 self-center"
-                    >
-                        {isExpanded ? (
-                            <ChevronDown className="w-4 h-4" />
-                        ) : (
-                            <ChevronRight className="w-4 h-4" />
-                        )}
-                    </Button>
-
-                    <div className="flex items-center justify-center p-2 w-24 bg-background-secondary border rounded-md mr-2">
-                        <ArrowRight className="w-4 h-4 mr-1" />
-                        <span className="text-sm font-medium">Step</span>
-                    </div>
-                </div>
-
-                <AutosizeTextarea
-                    value={step.step}
-                    onChange={(e) =>
-                        onUpdate({ ...step, step: e.target.value })
-                    }
-                    placeholder="Step"
-                    maxHeight={100}
-                    minHeight={0}
-                    className="flex-grow bg-background-secondary"
-                    onClick={(e) => e.stopPropagation()}
-                />
-
-                <Button
-                    variant="ghost"
-                    size="icon"
-                    onClick={(e) => {
-                        e.stopPropagation();
-                        onDelete();
-                    }}
-                    className="ml-2 min-w-fit"
-                >
-                    <Trash className="w-4 h-4" />
-                </Button>
-            </ClickableDiv>
-
-            {isExpanded && (
-                <div className="p-3 space-y-4 px-8">
-                    <Accordion type="multiple">
-                        <AccordionItem value="tools">
-                            <AccordionTrigger>
-                                <span className="flex items-center gap-2 justify-center">
-                                    <Wrench className="w-4 h-4" />
-                                    Tools
-                                </span>
-                            </AccordionTrigger>
-
-                            <AccordionContent className="p-1 pb-6">
-                                <BasicToolForm
-                                    value={step.tools}
-                                    onChange={(tools) =>
-                                        onUpdate({ ...step, tools })
-                                    }
-                                />
-                            </AccordionContent>
-                        </AccordionItem>
-
-                        <AccordionItem value="workflows">
-                            <AccordionTrigger>
-                                <span className="flex items-center gap-2 justify-center">
-                                    <Puzzle className="w-4 h-4" />
-                                    Workflows
-                                </span>
-                            </AccordionTrigger>
-
-                            <AccordionContent className="p-1 pb-6">
-                                <WorkflowSelectModule
-                                    onChange={(workflows) =>
-                                        onUpdate({ ...step, workflows })
-                                    }
-                                    selection={step.workflows || []}
-                                />
-                            </AccordionContent>
-                        </AccordionItem>
-
-                        <AccordionItem value="agents">
-                            <AccordionTrigger>
-                                <span className="flex items-center gap-2 justify-center">
-                                    <User className="w-4 h-4" />
-                                    Agents
-                                </span>
-                            </AccordionTrigger>
-
-                            <AccordionContent className="p-1 pb-6">
-                                <AgentSelectModule
-                                    onChange={(agents) =>
-                                        onUpdate({ ...step, agents })
-                                    }
-                                    selection={step.agents || []}
-                                />
-                            </AccordionContent>
-                        </AccordionItem>
-
-                        <AccordionItem value="advanced">
-                            <AccordionTrigger>
-                                <span className="flex items-center gap-2 justify-center">
-                                    <CogIcon className="w-4 h-4" />
-                                    Advanced
-                                </span>
-                            </AccordionTrigger>
-
-                            <AccordionContent className="p-1 pb-6 space-y-6">
-                                <div>
-                                    <label
-                                        htmlFor="temperature"
-                                        className="block text-sm font-medium text-gray-700 mb-1"
-                                    >
-                                        Temperature
-                                    </label>
-
-                                    <Input
-                                        id="temperature"
-                                        type="number"
-                                        value={step.temperature}
-                                        onChange={(e) =>
-                                            onUpdate({
-                                                ...step,
-                                                temperature: parseFloat(
-                                                    e.target.value
-                                                ),
-                                            })
-                                        }
-                                        placeholder="Temperature"
-                                        className="bg-background"
-                                    />
-                                </div>
-
-                                <div className="flex items-center space-x-2">
-                                    <Switch
-                                        checked={step.cache}
-                                        onCheckedChange={(checked) =>
-                                            onUpdate({
-                                                ...step,
-                                                cache: checked,
-                                            })
-                                        }
-                                    />
-
-                                    <span>Cache</span>
-                                </div>
-                            </AccordionContent>
-                        </AccordionItem>
-                    </Accordion>
-                </div>
-            )}
-        </div>
-    );
-}
diff --git a/ui/admin/app/components/workflow/steps/StepBase.tsx b/ui/admin/app/components/workflow/steps/StepBase.tsx
new file mode 100644
index 000000000..18a2d0b24
--- /dev/null
+++ b/ui/admin/app/components/workflow/steps/StepBase.tsx
@@ -0,0 +1,116 @@
+import { ChevronDown, ChevronRight, Trash } from "lucide-react";
+import { useState } from "react";
+
+import { Step, StepType, getDefaultStep } from "~/lib/model/workflows";
+import { cn } from "~/lib/utils";
+
+import { Button } from "~/components/ui/button";
+import { ClickableDiv } from "~/components/ui/clickable-div";
+import { AutosizeTextarea } from "~/components/ui/textarea";
+import { StepTypeSelect } from "~/components/workflow/steps/StepTypeSelect";
+
+export function StepBase({
+    className,
+    children,
+    step,
+    type,
+    onUpdate,
+    onDelete,
+}: {
+    className?: string;
+    children: React.ReactNode;
+    step: Step;
+    type: StepType;
+    onUpdate: (step: Step) => void;
+    onDelete: () => void;
+}) {
+    const [isExpanded, setIsExpanded] = useState(false);
+
+    const fieldConfig = getTextFieldConfig();
+
+    return (
+        <div className={cn("border rounded-md", className)}>
+            <ClickableDiv
+                className={cn(
+                    "flex items-start gap-2 p-3 bg-muted",
+                    isExpanded ? "rounded-t-md" : "rounded-md"
+                )}
+                onClick={() => setIsExpanded((prev) => !prev)}
+            >
+                <div className="flex items-center gap-2">
+                    <Button
+                        variant="ghost"
+                        size="icon"
+                        className="p-0 w-6 h-6 self-center"
+                    >
+                        {isExpanded ? (
+                            <ChevronDown className="w-4 h-4" />
+                        ) : (
+                            <ChevronRight className="w-4 h-4" />
+                        )}
+                    </Button>
+
+                    <StepTypeSelect
+                        value={type}
+                        onChange={(newType) => {
+                            if (newType !== type) {
+                                onUpdate(getDefaultStep(newType));
+                            }
+                        }}
+                    />
+                </div>
+
+                <AutosizeTextarea
+                    value={fieldConfig.value}
+                    onChange={(e) => fieldConfig.onChange(e.target.value)}
+                    placeholder={fieldConfig.placeholder}
+                    maxHeight={100}
+                    minHeight={0}
+                    className="flex-grow bg-background-secondary"
+                    onClick={(e) => e.stopPropagation()}
+                />
+
+                <Button
+                    variant="ghost"
+                    size="icon"
+                    onClick={(e) => {
+                        e.stopPropagation();
+                        onDelete();
+                    }}
+                >
+                    <Trash className="w-4 h-4" />
+                </Button>
+            </ClickableDiv>
+
+            {isExpanded && children}
+        </div>
+    );
+
+    function getTextFieldConfig() {
+        if (type === "if" && step.if) {
+            const copy = step.if;
+
+            return {
+                placeholder: "Condition",
+                value: step.if.condition,
+                onChange: (value: string) =>
+                    onUpdate({ ...step, if: { ...copy, condition: value } }),
+            };
+        } else if (type === "while" && step.while) {
+            const copy = step.while;
+
+            return {
+                placeholder: "Condition",
+                value: step.while.condition,
+                onChange: (value: string) =>
+                    onUpdate({ ...step, while: { ...copy, condition: value } }),
+            };
+        }
+
+        return {
+            placeholder: "Step",
+            value: step.step,
+            onChange: (value: string) => onUpdate({ ...step, step: value }),
+        };
+    }
+}
diff --git a/ui/admin/app/components/workflow/steps/StepContent.tsx b/ui/admin/app/components/workflow/steps/StepContent.tsx
new file mode 100644
index 000000000..89a148181
--- /dev/null
+++ b/ui/admin/app/components/workflow/steps/StepContent.tsx
@@ -0,0 +1,127 @@
+import { CogIcon, Puzzle, User, Wrench } from "lucide-react";
+
+import { Step } from "~/lib/model/workflows";
+
+import { AgentSelectModule } from "~/components/agent/shared/AgentSelect";
+import { BasicToolForm } from "~/components/tools/BasicToolForm";
+import {
+    Accordion,
+    AccordionContent,
+    AccordionItem,
+    AccordionTrigger,
+} from "~/components/ui/accordion";
+import { Input } from "~/components/ui/input";
+import { Switch } from "~/components/ui/switch";
+import { WorkflowSelectModule } from "~/components/workflow/WorkflowSelectModule";
+
+export function StepContent({
+    step,
+    onUpdate,
+}: {
+    step: Step;
+    onUpdate: (updatedStep: Step) => void;
+}) {
+    return (
+        <div className="p-3 space-y-4 px-8">
+            <Accordion type="multiple">
+                <AccordionItem value="tools">
+                    <AccordionTrigger>
+                        <span className="flex items-center gap-2 justify-center">
+                            <Wrench className="w-4 h-4" />
+                            Tools
+                        </span>
+                    </AccordionTrigger>
+
+                    <AccordionContent className="p-1 pb-6">
+                        <BasicToolForm
+                            value={step.tools}
+                            onChange={(tools) => onUpdate({ ...step, tools })}
+                        />
+                    </AccordionContent>
+                </AccordionItem>
+
+                <AccordionItem value="workflows">
+                    <AccordionTrigger>
+                        <span className="flex items-center gap-2 justify-center">
+                            <Puzzle className="w-4 h-4" />
+                            Workflows
+                        </span>
+                    </AccordionTrigger>
+
+                    <AccordionContent className="p-1 pb-6">
+                        <WorkflowSelectModule
+                            onChange={(workflows) =>
+                                onUpdate({ ...step, workflows })
+                            }
+                            selection={step.workflows || []}
+                        />
+                    </AccordionContent>
+                </AccordionItem>
+
+                <AccordionItem value="agents">
+                    <AccordionTrigger>
+                        <span className="flex items-center gap-2 justify-center">
+                            <User className="w-4 h-4" />
+                            Agents
+                        </span>
+                    </AccordionTrigger>
+
+                    <AccordionContent className="p-1 pb-6">
+                        <AgentSelectModule
+                            onChange={(agents) => onUpdate({ ...step, agents })}
+                            selection={step.agents || []}
+                        />
+                    </AccordionContent>
+                </AccordionItem>
+
+                <AccordionItem value="advanced">
+                    <AccordionTrigger>
+                        <span className="flex items-center gap-2 justify-center">
+                            <CogIcon className="w-4 h-4" />
+                            Advanced
+                        </span>
+                    </AccordionTrigger>
+
+                    <AccordionContent className="p-1 pb-6 space-y-6">
+                        <div>
+                            <label
+                                htmlFor="temperature"
+                                className="block text-sm font-medium text-gray-700 mb-1"
+                            >
+                                Temperature
+                            </label>
+
+                            <Input
+                                id="temperature"
+                                type="number"
+                                value={step.temperature}
+                                onChange={(e) =>
+                                    onUpdate({
+                                        ...step,
+                                        temperature: parseFloat(e.target.value),
+                                    })
+                                }
+                                placeholder="Temperature"
+                                className="bg-background"
+                            />
+                        </div>
+
+                        <div className="flex items-center space-x-2">
+                            <Switch
+                                checked={step.cache}
+                                onCheckedChange={(checked) =>
+                                    onUpdate({
+                                        ...step,
+                                        cache: checked,
+                                    })
+                                }
+                            />
+
+                            <span>Cache</span>
+                        </div>
+                    </AccordionContent>
+                </AccordionItem>
+            </Accordion>
+        </div>
+    );
+}
diff --git a/ui/admin/app/components/workflow/steps/StepRenderer.tsx b/ui/admin/app/components/workflow/steps/StepRenderer.tsx
index ae8a8f87c..975be532b 100644
--- a/ui/admin/app/components/workflow/steps/StepRenderer.tsx
+++ b/ui/admin/app/components/workflow/steps/StepRenderer.tsx
@@ -1,9 +1,10 @@
 import { Step } from "~/lib/model/workflows";
 
-import { IfComponent } from "~/components/workflow/steps/If";
-import { StepComponent } from "~/components/workflow/steps/Step";
+import { IfContent } from "~/components/workflow/steps/IfContent";
+import { StepBase } from "~/components/workflow/steps/StepBase";
+import { StepContent } from "~/components/workflow/steps/StepContent";
 import { TemplateComponent } from "~/components/workflow/steps/Template";
-import { WhileComponent } from "~/components/workflow/steps/While";
+import { WhileContent } from "~/components/workflow/steps/WhileContent";
 
 export function renderStep(
     step: Step,
@@ -21,37 +22,51 @@ export function renderStep(
         );
     } else if (step.if) {
         return (
-            <IfComponent
+            <StepBase
                 key={step.id}
-                ifCondition={step.if}
-                onUpdate={(updatedIf) => onUpdate({ ...step, if: updatedIf })}
+                step={step}
+                type="if"
+                onUpdate={onUpdate}
                 onDelete={onDelete}
-                renderStep={renderStep}
-                className="mb-4"
-            />
+            >
+                <IfContent
+                    ifCondition={step.if}
+                    onUpdate={(updatedIf) =>
+                        onUpdate({ ...step, if: updatedIf })
+                    }
+                    renderStep={renderStep}
+                />
+            </StepBase>
         );
     } else if (step.while) {
         return (
-            <WhileComponent
+            <StepBase
                 key={step.id}
-                whileCondition={step.while}
-                onUpdate={(updatedWhile) =>
-                    onUpdate({ ...step, while: updatedWhile })
-                }
+                step={step}
+                type="while"
+                onUpdate={onUpdate}
                 onDelete={onDelete}
-                renderStep={renderStep}
-                className="mb-4"
-            />
+            >
+                <WhileContent
+                    whileCondition={step.while}
+                    onUpdate={(updatedWhile) =>
+                        onUpdate({ ...step, while: updatedWhile })
+                    }
+                    renderStep={renderStep}
+                />
+            </StepBase>
         );
     } else {
         return (
-            <StepComponent
+            <StepBase
                 key={step.id}
                 step={step}
+                type="command"
                 onUpdate={onUpdate}
                 onDelete={onDelete}
-                className="mb-4"
-            />
+            >
+                <StepContent key={step.id} step={step} onUpdate={onUpdate} />
+            </StepBase>
         );
     }
 }
diff --git a/ui/admin/app/components/workflow/steps/StepTypeSelect.tsx b/ui/admin/app/components/workflow/steps/StepTypeSelect.tsx
new file mode 100644
index 000000000..433b92c4c
--- /dev/null
+++ b/ui/admin/app/components/workflow/steps/StepTypeSelect.tsx
@@ -0,0 +1,65 @@
+import { ArrowRight, RotateCwIcon, SplitIcon } from "lucide-react";
+
+import { StepType } from "~/lib/model/workflows";
+
+import { ClickableDiv } from "~/components/ui/clickable-div";
+import {
+    Select,
+    SelectContent,
+    SelectItem,
+    SelectTrigger,
+    SelectValue,
+} from "~/components/ui/select";
+
+type StepSelectProps = {
+    value: StepType;
+    onChange: (type: StepType) => void;
+};
+
+export function StepTypeSelect({ value, onChange }: StepSelectProps) {
+    const options = Object.values(StepType).map((type) => ({
+        id: type,
+        name: StepLabelMap[type],
+    }));
+
+    return (
+        <ClickableDiv onClick={(e) => e.stopPropagation()}>
+            <Select value={value} onValueChange={onChange}>
+                <SelectTrigger className="bg-background w-28">
+                    <SelectValue />
+                </SelectTrigger>
+
+                <SelectContent>
+                    {options.map((option) => (
+                        <SelectItem key={option.id} value={option.id}>
+                            <div className="flex items-center gap-2">
+                                <StepTypeIcon type={option.id} />
+                                {option.name}
+                            </div>
+                        </SelectItem>
+                    ))}
+                </SelectContent>
+            </Select>
+        </ClickableDiv>
+    );
+}
+
+const StepLabelMap: Record<StepType, string> = {
+    [StepType.Command]: "Step",
+    [StepType.If]: "If",
+    [StepType.While]: "While",
+    // [StepType.Template]: "Template",
+};
+
+const IconMap: Record<StepType, React.ElementType> = {
+    [StepType.Command]: ArrowRight,
+    [StepType.If]: SplitIcon,
+    [StepType.While]: RotateCwIcon,
+    // [StepType.Template]: PuzzleIcon,
+};
+
+function StepTypeIcon({ type }: { type: StepType }) {
+    const Icon = IconMap[type] ?? ArrowRight;
+
+    return <Icon className="w-4 h-4" />;
+}
diff --git a/ui/admin/app/components/workflow/steps/While.tsx b/ui/admin/app/components/workflow/steps/While.tsx
deleted file mode 100644
index 090fa2089..000000000
--- a/ui/admin/app/components/workflow/steps/While.tsx
+++ /dev/null
@@ -1,153 +0,0 @@
-import { ChevronDown, ChevronRight, RotateCw, Trash } from "lucide-react";
-import { useState } from "react";
-
-import { Step, While } from "~/lib/model/workflows";
-import { cn } from "~/lib/utils";
-
-import { Button } from "~/components/ui/button";
-import { ClickableDiv } from "~/components/ui/clickable-div";
-import { Input } from "~/components/ui/input";
-import { AutosizeTextarea } from "~/components/ui/textarea";
-import { AddStepButton } from "~/components/workflow/steps/AddStep";
-
-export function WhileComponent({
-    whileCondition,
-    onUpdate,
-    onDelete,
-    renderStep,
-    className,
-}: {
-    whileCondition: While;
-    onUpdate: (updatedWhile: While) => void;
-    onDelete: () => void;
-    renderStep: (
-        step: Step,
-        onUpdate: (updatedStep: Step) => void,
-        onDelete: () => void
-    ) => React.ReactNode;
-    className?: string;
-}) {
-    const [isExpanded, setIsExpanded] = useState(false);
-
-    const addStep = (newStep: Step) => {
-        onUpdate({
-            ...whileCondition,
-            steps: [...(whileCondition.steps || []), newStep],
-        });
-    };
-
-    const updateNestedStep = (index: number, updatedStep: Step) => {
-        const newSteps = [...(whileCondition.steps || [])];
-        newSteps[index] = updatedStep;
-        onUpdate({ ...whileCondition, steps: newSteps });
-    };
-
-    const deleteNestedStep = (index: number) => {
-        const newSteps = (whileCondition.steps || []).filter(
-            (_, i) => i !== index
-        );
-        onUpdate({ ...whileCondition, steps: newSteps });
-    };
-
-    return (
-        <div className={cn("border rounded-md", className)}>
-            <ClickableDiv
-                className={cn(
-                    "flex items-start p-3 bg-secondary",
-                    isExpanded ? "rounded-t-md" : "rounded-md"
-                )}
-                onClick={() => setIsExpanded((prev) => !prev)}
-            >
-                <div className="flex items-center">
-                    <Button
-                        variant="ghost"
-                        size="icon"
-                        className="p-0 w-6 h-6 mr-2 self-center"
-                        onClick={() => setIsExpanded(!isExpanded)}
-                    >
-                        {isExpanded ? (
-                            <ChevronDown className="w-4 h-4" />
-                        ) : (
-                            <ChevronRight className="w-4 h-4" />
-                        )}
-                    </Button>
-
-                    <div className="flex items-center justify-center w-24 p-2 border bg-background rounded-md mr-2">
-                        <RotateCw className="w-4 h-4 mr-1" />
-                        <span className="text-sm font-medium">While</span>
-                    </div>
-                </div>
-
-                <AutosizeTextarea
-                    value={whileCondition.condition}
-                    onChange={(e) =>
-                        onUpdate({
-                            ...whileCondition,
-                            condition: e.target.value,
-                        })
-                    }
-                    minHeight={0}
-                    maxHeight={100}
-                    placeholder="Condition"
-                    className="flex-grow bg-background"
-                    onClick={(e) => e.stopPropagation()}
-                />
-                <Button
-                    variant="ghost"
-                    size="icon"
-                    onClick={(e) => {
-                        e.stopPropagation();
-                        onDelete();
-                    }}
-                    className="ml-2"
-                >
-                    <Trash className="w-4 h-4" />
-                </Button>
-            </ClickableDiv>
-
-            {isExpanded && (
-                <div className="p-3 space-y-4">
-                    <div>
-                        <label
-                            htmlFor="maxLoops"
-                            className="block text-sm font-medium text-gray-700 mb-1"
-                        >
-                            Max Loops
-                        </label>
-                        <Input
-                            id="maxLoops"
-                            type="number"
-                            value={whileCondition.maxLoops}
-                            onChange={(e) =>
-                                onUpdate({
-                                    ...whileCondition,
-                                    maxLoops: parseInt(e.target.value),
-                                })
-                            }
-                            placeholder="Max Loops"
-                            className="bg-background"
-                        />
-                    </div>
-
-                    <div className="space-y-2">
-                        <h4 className="font-semibold">Steps:</h4>
-                        {whileCondition.steps?.map((step, index) => (
-                            <div key={index} className="ml-4">
-                                {renderStep(
-                                    step,
-                                    (updatedStep) =>
-                                        updateNestedStep(index, updatedStep),
-                                    () => deleteNestedStep(index)
-                                )}
-                            </div>
-                        ))}
-
-                        <div className="ml-4">
-                            <AddStepButton onAddStep={addStep} />
-                        </div>
-                    </div>
-                </div>
-            )}
-        </div>
-    );
-}
diff --git a/ui/admin/app/components/workflow/steps/WhileContent.tsx b/ui/admin/app/components/workflow/steps/WhileContent.tsx
new file mode 100644
index 000000000..aa3086929
--- /dev/null
+++ b/ui/admin/app/components/workflow/steps/WhileContent.tsx
@@ -0,0 +1,83 @@
+import { Step, While } from "~/lib/model/workflows";
+
+import { Input } from "~/components/ui/input";
+import { AddStepButton } from "~/components/workflow/steps/AddStep";
+
+export function WhileContent({
+    whileCondition,
+    onUpdate,
+    renderStep,
+}: {
+    whileCondition: While;
+    onUpdate: (updatedWhile: While) => void;
+    renderStep: (
+        step: Step,
+        onUpdate: (updatedStep: Step) => void,
+        onDelete: () => void
+    ) => React.ReactNode;
+    className?: string;
+}) {
+    const addStep = (newStep: Step) => {
+        onUpdate({
+            ...whileCondition,
+            steps: [...(whileCondition.steps || []), newStep],
+        });
+    };
+
+    const updateNestedStep = (index: number, updatedStep: Step) => {
+        const newSteps = [...(whileCondition.steps || [])];
+        newSteps[index] = updatedStep;
+        onUpdate({ ...whileCondition, steps: newSteps });
+    };
+
+    const deleteNestedStep = (index: number) => {
+        const newSteps = (whileCondition.steps || []).filter(
+            (_, i) => i !== index
+        );
+        onUpdate({ ...whileCondition, steps: newSteps });
+    };
+
+    return (
+        <div className="p-3 space-y-4">
+            <div>
+                <label
+                    htmlFor="maxLoops"
+                    className="block text-sm font-medium text-gray-700 mb-1"
+                >
+                    Max Loops
+                </label>
+                <Input
+                    id="maxLoops"
+                    type="number"
+                    value={whileCondition.maxLoops}
+                    onChange={(e) =>
+                        onUpdate({
+                            ...whileCondition,
+                            maxLoops: parseInt(e.target.value),
+                        })
+                    }
+                    placeholder="Max Loops"
+                    className="bg-background"
+                />
+            </div>
+
+            <div className="space-y-2">
+                <h4 className="font-semibold">Steps:</h4>
+                {whileCondition.steps?.map((step, index) => (
+                    <div key={index} className="ml-4">
+                        {renderStep(
+                            step,
+                            (updatedStep) =>
+                                updateNestedStep(index, updatedStep),
+                            () => deleteNestedStep(index)
+                        )}
+                    </div>
+                ))}
+
+                <div className="ml-4">
+                    <AddStepButton onAddStep={addStep} />
+                </div>
+            </div>
+        </div>
+    );
+}
diff --git a/ui/admin/app/lib/model/workflows.ts b/ui/admin/app/lib/model/workflows.ts
index 19d8f2d75..79e8f043d 100644
--- a/ui/admin/app/lib/model/workflows.ts
+++ b/ui/admin/app/lib/model/workflows.ts
@@ -6,6 +6,14 @@ export type WorkflowBase = AgentBase & {
     output: string;
 };
 
+export const StepType = {
+    Command: "command",
+    If: "if",
+    While: "while",
+    // Template: "template",
+} as const;
+export type StepType = (typeof StepType)[keyof typeof StepType];
+
 export type Step = {
     id: string;
     name: string;
@@ -50,14 +58,24 @@ export type Workflow = EntityMeta &
 export type CreateWorkflow = Partial<WorkflowBase> & Pick<WorkflowBase, "name">;
 export type UpdateWorkflow = WorkflowBase;
 
-export const getDefaultStep = (): Step => ({
-    id: crypto.randomUUID(),
-    name: "",
-    description: "",
-    step: "",
-    cache: false,
-    temperature: 0,
-    tools: [],
-    agents: [],
-    workflows: [],
-});
+export const getDefaultStep = (type: StepType = StepType.Command): Step => {
+    const newStep: Step = {
+        id: crypto.randomUUID(),
+        name: "",
+        description: "",
+        step: "",
+        cache: false,
+        temperature: 0,
+        tools: [],
+        agents: [],
+        workflows: [],
+    };
+
+    if (type === StepType.If) {
+        newStep.if = { condition: "", steps: [], else: [] };
+    } else if (type === StepType.While) {
+        newStep.while = { condition: "", maxLoops: 0, steps: [] };
+    }
+
+    return newStep;
+};