From 3f850bbd7f32ef8336f03d6e32aa4f8ebc2056b6 Mon Sep 17 00:00:00 2001
From: Kyle Corbitt <kyle@corbt.com>
Date: Mon, 26 Jun 2023 14:32:45 -0700
Subject: [PATCH] editing template variables works

---
 .../OutputsTable/NewScenarioButton.tsx        |   1 -
 .../OutputsTable/NewVariantButton.tsx         |   1 -
 .../OutputsTable/ScenarioEditor.tsx           |   4 +-
 .../OutputsTable/VariantConfigEditor.tsx      |   3 +-
 src/components/OutputsTable/VariantHeader.tsx |   2 +-
 src/components/OutputsTable/index.tsx         |  10 +-
 src/server/ScenarioHeader.tsx                 | 121 ++++++++++++++++++
 src/server/api/routers/experiments.router.ts  |  28 ++++
 src/utils/theme.ts                            |  17 +++
 9 files changed, 173 insertions(+), 14 deletions(-)
 create mode 100644 src/server/ScenarioHeader.tsx

diff --git a/src/components/OutputsTable/NewScenarioButton.tsx b/src/components/OutputsTable/NewScenarioButton.tsx
index 68801cb..43827ed 100644
--- a/src/components/OutputsTable/NewScenarioButton.tsx
+++ b/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"
diff --git a/src/components/OutputsTable/NewVariantButton.tsx b/src/components/OutputsTable/NewVariantButton.tsx
index 02cf232..cde24ad 100644
--- a/src/components/OutputsTable/NewVariantButton.tsx
+++ b/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"
diff --git a/src/components/OutputsTable/ScenarioEditor.tsx b/src/components/OutputsTable/ScenarioEditor.tsx
index 3ae4a9e..692f684 100644
--- a/src/components/OutputsTable/ScenarioEditor.tsx
+++ b/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>
diff --git a/src/components/OutputsTable/VariantConfigEditor.tsx b/src/components/OutputsTable/VariantConfigEditor.tsx
index 8397ef3..a3a0044 100644
--- a/src/components/OutputsTable/VariantConfigEditor.tsx
+++ b/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>
diff --git a/src/components/OutputsTable/VariantHeader.tsx b/src/components/OutputsTable/VariantHeader.tsx
index f3a4ca6..86d97bd 100644
--- a/src/components/OutputsTable/VariantHeader.tsx
+++ b/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>
diff --git a/src/components/OutputsTable/index.tsx b/src/components/OutputsTable/index.tsx
index 2d0a4d6..f1af906 100644
--- a/src/components/OutputsTable/index.tsx
+++ b/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}>
-          <Heading size="sm" fontWeight="bold">
-            Scenario
-          </Heading>
-        </Box>
+        <ScenarioHeader />
       </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>
       ))}
diff --git a/src/server/ScenarioHeader.tsx b/src/server/ScenarioHeader.tsx
new file mode 100644
index 0000000..2eca6fa
--- /dev/null
+++ b/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>
+  );
+}
diff --git a/src/server/api/routers/experiments.router.ts b/src/server/api/routers/experiments.router.ts
index c7e56aa..f618282 100644
--- a/src/server/api/routers/experiments.router.ts
+++ b/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;
+    }),
 });
diff --git a/src/utils/theme.ts b/src/utils/theme.ts
index 41671d0..083d3b3 100644
--- a/src/utils/theme.ts
+++ b/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;