editing template variables works

src/components/OutputsTable/NewScenarioButton.tsx

@@ -18,7 +18,6 @@ export default function NewScenarioButton() {
   return (
     <Button
       w="100%"
-      borderRadius={0}
       alignItems="center"
       justifyContent="flex-start"
       fontWeight="normal"

src/components/OutputsTable/NewVariantButton.tsx

@@ -20,7 +20,6 @@ export default function NewVariantButton() {
   return (
     <Button
       w="100%"
-      borderRadius={0}
       alignItems="center"
       justifyContent="center"
       fontWeight="normal"

src/components/OutputsTable/ScenarioEditor.tsx

@@ -128,7 +128,6 @@ export default function ScenarioEditor({
                 {key}
               </Box>
               <Textarea
-                borderRadius={0}
                 px={2}
                 py={1}
                 placeholder="empty"
@@ -153,7 +152,6 @@ export default function ScenarioEditor({
           <HStack justify="right">
             <Button
               size="sm"
-              borderRadius={0}
               onMouseDown={() => {
                 setValues(savedValues);
               }}
@@ -161,7 +159,7 @@ export default function ScenarioEditor({
             >
               Reset
             </Button>
-            <Button size="sm" borderRadius={0} onMouseDown={onSave} colorScheme="blue">
+            <Button size="sm" onMouseDown={onSave} colorScheme="blue">
               Save
             </Button>
           </HStack>

src/components/OutputsTable/VariantConfigEditor.tsx

@@ -165,12 +165,11 @@ export default function VariantConfigEditor(props: { variant: PromptVariant }) {
               editorRef.current?.setValue(savedConfig);
               checkForChanges();
             }}
-            borderRadius={0}
           >
             Reset
           </Button>
           <Tooltip label={`${modifierKey} + Enter`}>
-            <Button size="sm" onClick={onSave} colorScheme="blue" borderRadius={0}>
+            <Button size="sm" onClick={onSave} colorScheme="blue">
               Save
             </Button>
           </Tooltip>

src/components/OutputsTable/VariantHeader.tsx

@@ -91,7 +91,7 @@ export default function VariantHeader(props: { variant: PromptVariant }) {
         px={cellPadding.x}
       />
       <Tooltip label="Hide Variant" hasArrow>
-        <Button variant="ghost" colorScheme="gray" size="sm" onClick={onHide} borderRadius={0}>
+        <Button variant="ghost" colorScheme="gray" size="sm" onClick={onHide}>
           <Icon as={BsX} boxSize={6} />
         </Button>
       </Tooltip>

src/components/OutputsTable/index.tsx

@@ -9,6 +9,7 @@ import VariantConfigEditor from "./VariantConfigEditor";
 import VariantHeader from "./VariantHeader";
 import type { Scenario, PromptVariant } from "./types";
 import { cellPadding } from "../constants";
+import ScenarioHeader from "~/server/ScenarioHeader";
 
 const stickyHeaderStyle: SystemStyleObject = {
   position: "sticky",
@@ -30,6 +31,7 @@ const ScenarioRow = (props: { scenario: Scenario; variants: PromptVariant[] }) =
         onMouseEnter={() => setIsHovered(true)}
         onMouseLeave={() => setIsHovered(false)}
         sx={isHovered ? highlightStyle : undefined}
+        borderLeftWidth={1}
       >
         <ScenarioEditor scenario={props.scenario} hovered={isHovered} />
       </GridItem>
@@ -77,14 +79,10 @@ export default function OutputsTable({ experimentId }: { experimentId: string |
       }}
     >
       <GridItem display="flex" alignItems="flex-end" rowSpan={2}>
-        <Box sx={stickyHeaderStyle} flex={1} px={cellPadding.x} py={cellPadding.y}>
+        <ScenarioHeader />
-          <Heading size="sm" fontWeight="bold">
-            Scenario
-          </Heading>
-        </Box>
       </GridItem>
       {variants.data.map((variant) => (
-        <GridItem key={variant.uiId} padding={0} sx={stickyHeaderStyle}>
+        <GridItem key={variant.uiId} padding={0} sx={stickyHeaderStyle} borderTopWidth={1}>
           <VariantHeader variant={variant} />
         </GridItem>
       ))}

src/server/ScenarioHeader.tsx

@@ -0,0 +1,121 @@
+import { Text, Box, Button, HStack, Heading, Icon, Input, Stack, Code } from "@chakra-ui/react";
+import { useState } from "react";
+import { BsCheck, BsChevronDown, BsX } from "react-icons/bs";
+import { cellPadding } from "~/components/constants";
+import { api } from "~/utils/api";
+import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
+
+export default function ScenarioHeader() {
+  const experiment = useExperiment();
+
+  const initialVariables = experiment.data?.TemplateVariable ?? [];
+
+  const [variables, setVariables] = useState<string[]>(initialVariables.map((v) => v.label));
+  const [newVariable, setNewVariable] = useState<string>("");
+
+  const [editing, setEditing] = useState(true);
+
+  const utils = api.useContext();
+  const setVarsMutation = api.experiments.setTemplateVariables.useMutation();
+  const [onSave] = useHandledAsyncCallback(async () => {
+    if (!experiment.data?.id) return;
+    setEditing(false);
+    await setVarsMutation.mutateAsync({
+      id: experiment.data.id,
+      labels: variables,
+    });
+    await utils.experiments.get.invalidate();
+  }, [setVarsMutation, experiment.data?.id, variables]);
+
+  return (
+    <Stack flex={1} px={cellPadding.x} py={cellPadding.y}>
+      <HStack>
+        <Heading size="sm" fontWeight="bold" flex={1}>
+          Scenario
+        </Heading>
+        {editing ? (
+          <HStack>
+            <Button size="xs" colorScheme="gray" onClick={() => setEditing(false)}>
+              Cancel
+            </Button>
+            <Button size="xs" colorScheme="blue" onClick={onSave}>
+              Save
+            </Button>
+          </HStack>
+        ) : (
+          <Button
+            size="xs"
+            variant="outline"
+            colorScheme="blue"
+            rightIcon={<Icon as={BsChevronDown} />}
+            onClick={() => setEditing(true)}
+          >
+            Edit Vars
+          </Button>
+        )}
+      </HStack>
+      {editing && (
+        <Stack spacing={2} pt={2}>
+          <Text fontSize="sm">
+            You can use variables in your prompt text inside <Code>{`{{curly_braces}}`}</Code>.
+          </Text>
+          <HStack spacing={0}>
+            <Input
+              placeholder="Add Variable Name"
+              size="sm"
+              borderTopRadius={0}
+              borderRightRadius={0}
+              value={newVariable}
+              onChange={(e) => setNewVariable(e.target.value)}
+            />
+            <Button
+              size="xs"
+              height="100%"
+              borderLeftRadius={0}
+              isDisabled={newVariable.length === 0 || variables.includes(newVariable)}
+              onClick={() => {
+                setVariables([...variables, newVariable]);
+                setNewVariable("");
+              }}
+            >
+              <Icon as={BsCheck} boxSize={8} />
+            </Button>
+          </HStack>
+
+          <HStack spacing={2} py={4} wrap="wrap">
+            {variables.map((label, i) => (
+              <HStack
+                key={label}
+                spacing={0}
+                bgColor="blue.100"
+                color="blue.600"
+                pl={2}
+                pr={0}
+                fontWeight="bold"
+              >
+                <Text fontSize="sm" flex={1}>
+                  {label}
+                </Text>
+                <Button
+                  size="xs"
+                  variant="ghost"
+                  colorScheme="blue"
+                  p="unset"
+                  minW="unset"
+                  px="unset"
+                  onClick={() => {
+                    const newVariables = [...variables];
+                    newVariables.splice(i, 1);
+                    setVariables(newVariables);
+                  }}
+                >
+                  <Icon as={BsX} boxSize={6} color="blue.800" />
+                </Button>
+              </HStack>
+            ))}
+          </HStack>
+        </Stack>
+      )}
+    </Stack>
+  );
+}

src/server/api/routers/experiments.router.ts

@@ -21,4 +21,32 @@ export const experimentsRouter = createTRPCRouter({
       },
     });
   }),
+
+  setTemplateVariables: publicProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        labels: z.array(z.string()),
+      })
+    )
+    .mutation(async ({ input }) => {
+      const existing = await prisma.templateVariable.findMany({
+        where: { experimentId: input.id },
+      });
+      const toDelete = existing.filter((e) => !input.labels.includes(e.label));
+
+      const toCreate = new Set(
+        input.labels.filter((l) => !existing.map((e) => e.label).includes(l))
+      ).values();
+
+      await prisma.$transaction([
+        prisma.templateVariable.deleteMany({
+          where: { id: { in: toDelete.map((e) => e.id) } },
+        }),
+        prisma.templateVariable.createMany({
+          data: [...toCreate].map((l) => ({ label: l, experimentId: input.id })),
+        }),
+      ]);
+      return null;
+    }),
 });

src/utils/theme.ts

@@ -8,6 +8,23 @@ const theme = extendTheme({
     heading: systemFont,
     body: systemFont,
   },
+
+  components: {
+    Button: {
+      baseStyle: {
+        borderRadius: "sm",
+      },
+    },
+    Input: {
+      sizes: {
+        md: {
+          field: {
+            borderRadius: "sm",
+          },
+        },
+      },
+    },
+  },
 });
 
 export default theme;