feat: add enum to UI (#292)

ucbepic · Jan 25, 2025 · 8b33bc9 · 8b33bc9
1 parent 2a259a0
commit 8b33bc9
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 9 deletions.
diff --git a/website/src/app/api/utils.ts b/website/src/app/api/utils.ts
@@ -158,6 +158,13 @@ export function generatePipelineConfig(
           return `{${Object.entries(subSchema)
             .map(([k, v]) => `${k}: ${v}`)
             .join(", ")}}`;
+        } else if (item.type === "enum") {
+          if (!item.enumValues?.length) {
+            throw new Error(
+              `Enum type must specify its values for field: ${item.key}`
+            );
+          }
+          return `enum[${item.enumValues.join(", ")}]`;
         } else {
           return item.type;
         }

diff --git a/website/src/app/types.ts b/website/src/app/types.ts
@@ -42,12 +42,14 @@ export type SchemaType =
   | "int"
   | "boolean"
   | "list"
-  | "dict";
+  | "dict"
+  | "enum";
 
 export interface SchemaItem {
   key: string;
   type: SchemaType;
   subType?: SchemaItem | SchemaItem[];
+  enumValues?: string[];
 }
 
 export interface UserNote {

diff --git a/website/src/components/operations/args.tsx b/website/src/components/operations/args.tsx
@@ -126,6 +126,7 @@ export const SchemaForm: React.FC<SchemaFormProps> = React.memo(
                       : value === "dict"
                       ? [{ key: "", type: "string" }]
                       : undefined,
+                  enumValues: value === "enum" ? ["", ""] : undefined,
                 });
               }}
             >
@@ -139,6 +140,7 @@ export const SchemaForm: React.FC<SchemaFormProps> = React.memo(
                 <SelectItem value="boolean">boolean</SelectItem>
                 <SelectItem value="list">list</SelectItem>
                 <SelectItem value="dict">dict</SelectItem>
+                <SelectItem value="enum">enum</SelectItem>
               </SelectContent>
             </Select>
             {!isList && (
@@ -151,6 +153,79 @@ export const SchemaForm: React.FC<SchemaFormProps> = React.memo(
                 <Trash2 size={16} />
               </Button>
             )}
+            {item.type === "enum" && (
+              <div className="w-full mt-1 ml-4">
+                <Label className="text-sm text-gray-500 mb-1">
+                  Enum Values
+                </Label>
+                <div className="flex items-center gap-2">
+                  <div className="flex-1 overflow-x-auto flex items-center gap-2 pb-2">
+                    {(item.enumValues || ["", ""]).map((value, enumIndex) => (
+                      <div
+                        key={enumIndex}
+                        className="flex-none flex items-center gap-1"
+                      >
+                        <Input
+                          value={value}
+                          onChange={(e) => {
+                            const newValues = [
+                              ...(item.enumValues || ["", ""]),
+                            ];
+                            newValues[enumIndex] = e.target.value;
+                            updateItem(index, {
+                              ...item,
+                              enumValues: newValues,
+                            });
+                          }}
+                          placeholder={`value${enumIndex + 1}`}
+                          className={`w-32 ${!value ? "border-red-500" : ""}`}
+                        />
+                        {item.enumValues && item.enumValues.length > 2 && (
+                          <Button
+                            variant="ghost"
+                            size="sm"
+                            onClick={() => {
+                              const newValues = item.enumValues?.filter(
+                                (_, i) => i !== enumIndex
+                              );
+                              updateItem(index, {
+                                ...item,
+                                enumValues: newValues,
+                              });
+                            }}
+                            className="p-1 h-7 w-7 hover:bg-destructive/10"
+                          >
+                            <Trash2 size={14} className="text-destructive" />
+                          </Button>
+                        )}
+                      </div>
+                    ))}
+                  </div>
+                  <Button
+                    variant="outline"
+                    size="sm"
+                    onClick={() => {
+                      const newValues = [...(item.enumValues || ["", ""]), ""];
+                      updateItem(index, {
+                        ...item,
+                        enumValues: newValues,
+                      });
+                    }}
+                    className="flex-none"
+                  >
+                    <Plus size={16} className="mr-1" />
+                    Add Value
+                  </Button>
+                </div>
+                {(!item.enumValues?.length ||
+                  item.enumValues.length < 2 ||
+                  item.enumValues.some((v) => !v)) && (
+                  <div className="text-red-500 text-sm mt-1">
+                    At least two non-empty enum values are required
+                  </div>
+                )}
+              </div>
+            )}
             {item.type === "list" && item.subType && (
               <div className="w-full mt-1 ml-4 flex items-center">
                 <span className="mr-2 text-sm text-gray-500">List type:</span>

diff --git a/website/src/components/utils.ts b/website/src/components/utils.ts
@@ -12,6 +12,23 @@ export const schemaDictToItemSet = (
           type: "list",
           subType: { key: "0", type: subType as SchemaType },
         };
+      } else if (type.startsWith("enum[") && type.endsWith("]")) {
+        const enumValuesStr = type.slice(5, -1);
+        const enumValues = enumValuesStr
+          .split(",")
+          .map((v) => v.trim())
+          .filter((v) => v);
+        if (enumValues.length < 2) {
+          console.error(
+            `Invalid enum values for ${key}: ${type}. Using string type instead.`
+          );
+          return { key, type: "string" };
+        }
+        return {
+          key,
+          type: "enum",
+          enumValues,
+        };
       } else if (type.startsWith("{") && type.endsWith("}")) {
         try {
           const subSchema = JSON.parse(type);

diff --git a/website/src/hooks/useRestorePipeline.ts b/website/src/hooks/useRestorePipeline.ts
@@ -1,7 +1,6 @@
 import { useCallback } from "react";
 import yaml from "js-yaml";
 import { v4 as uuidv4 } from "uuid";
-import path from "path";
 import { Operation, File } from "@/app/types";
 import { schemaDictToItemSet } from "@/components/utils";
 import { useToast } from "@/hooks/use-toast";
@@ -21,6 +20,7 @@ interface YAMLOperation {
   };
   validate?: unknown;
   sample?: unknown;
+  reduce_key?: string | string[];
   [key: string]: unknown;
 }
 
@@ -86,16 +86,31 @@ export const useRestorePipeline = ({
                     output,
                     validate,
                     sample,
+                    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+                    reduce_key, // Explicitly ignore reduce_key as it's handled in otherKwargs
                     ...otherKwargs
                   } = op;
 
+                  // Convert all otherKwargs values to strings
+                  const stringifiedKwargs: Record<string, string> =
+                    Object.entries(otherKwargs).reduce(
+                      (acc, [key, value]) => ({
+                        ...acc,
+                        [key]:
+                          typeof value === "object"
+                            ? JSON.stringify(value)
+                            : String(value),
+                      }),
+                      {} as Record<string, string>
+                    );
+
                   // If the operation type is 'reduce', ensure reduce_key is a list
-                  if (type === "reduce" && otherKwargs.reduce_key) {
-                    otherKwargs.reduce_key = Array.isArray(
-                      otherKwargs.reduce_key
-                    )
-                      ? otherKwargs.reduce_key
-                      : [otherKwargs.reduce_key];
+                  if (type === "reduce" && op.reduce_key) {
+                    stringifiedKwargs.reduce_key = JSON.stringify(
+                      Array.isArray(op.reduce_key)
+                        ? op.reduce_key
+                        : [op.reduce_key]
+                    );
                   }
 
                   return {
@@ -120,7 +135,7 @@ export const useRestorePipeline = ({
                       : undefined,
                     validate,
                     sample,
-                    otherKwargs,
+                    otherKwargs: stringifiedKwargs,
                     visibility: true,
                   } as Operation;
                 })