Allow user to create a version of their current prompt with a new model (#58)

* Add dropdown header for model switching * Allow variant duplication * Fix prettier * Use env variable to restrict prisma logs * Fix env.mjs * Remove unnecessary scroll bar from function call output * Properly record when 404 error occurs in queryLLM task * Add SelectedModelInfo in SelectModelModal * Add react-select * Calculate new prompt after switching model * Send newly selected model with creation request * Get new prompt construction function back from GPT-4 * Fix prettier * Fix prettier
2023-07-18 18:24:04 -07:00
parent fa5b1ab1c5
commit e0e64c4207
18 changed files with 634 additions and 113 deletions
--- a/src/components/OutputsTable/OutputCell/OutputCell.tsx
+++ b/src/components/OutputsTable/OutputCell/OutputCell.tsx
@@ -106,7 +106,7 @@ export default function OutputCell({
        h="100%"
        fontSize="xs"
        flexWrap="wrap"
-        overflowX="auto"
+        overflowX="hidden"
        justifyContent="space-between"
      >
        <VStack w="full" flex={1} spacing={0}>
--- a/src/components/OutputsTable/index.tsx
+++ b/src/components/OutputsTable/index.tsx
@@ -4,7 +4,7 @@ import NewScenarioButton from "./NewScenarioButton";
 import NewVariantButton from "./NewVariantButton";
 import ScenarioRow from "./ScenarioRow";
 import VariantEditor from "./VariantEditor";
-import VariantHeader from "./VariantHeader";
+import VariantHeader from "../VariantHeader/VariantHeader";
 import VariantStats from "./VariantStats";
 import { ScenariosHeader } from "./ScenariosHeader";
 import { stickyHeaderStyle } from "./styles";
--- a/src/components/SelectModelModal/ModelStatsCard.tsx
+++ b/src/components/SelectModelModal/ModelStatsCard.tsx
@@ -0,0 +1,83 @@
+import {
+  Heading,
+  VStack,
+  Text,
+  HStack,
+  type StackProps,
+  Icon,
+  Button,
+  GridItem,
+  SimpleGrid,
+} from "@chakra-ui/react";
+import { BsChevronRight } from "react-icons/bs";
+import { modelStats } from "~/server/modelStats";
+import { type SupportedModel } from "~/server/types";
+
+export const ModelStatsCard = ({ label, model }: { label: string; model: SupportedModel }) => {
+  const stats = modelStats[model];
+  return (
+    <VStack w="full" spacing={6} bgColor="gray.100" p={4} borderRadius={8}>
+      <HStack w="full" justifyContent="space-between">
+        <Text fontWeight="bold" fontSize="xs">
+          {label}
+        </Text>
+        <Button variant="link" onClick={() => window.open(stats.learnMoreUrl, "_blank")}>
+          <HStack alignItems="center" spacing={0} color="blue.500" fontWeight="bold">
+            <Text fontSize="xs">Learn More</Text>
+            <Icon as={BsChevronRight} boxSize={3} strokeWidth={1} />
+          </HStack>
+        </Button>
+      </HStack>
+      <HStack w="full" justifyContent="space-between">
+        <Heading as="h3" size="md">
+          {model}
+        </Heading>
+        <Text fontWeight="bold">{stats.provider}</Text>
+      </HStack>
+      <SimpleGrid
+        w="full"
+        justifyContent="space-between"
+        alignItems="flex-start"
+        fontSize="sm"
+        columns={{ base: 2, md: 4 }}
+      >
+        <SelectedModelLabeledInfo label="Context" info={stats.contextLength} />
+        <SelectedModelLabeledInfo
+          label="Input"
+          info={
+            <Text>
+              ${(stats.promptTokenPrice * 1000).toFixed(3)}
+              <Text color="gray.500"> / 1K tokens</Text>
+            </Text>
+          }
+        />
+        <SelectedModelLabeledInfo
+          label="Output"
+          info={
+            <Text>
+              ${(stats.promptTokenPrice * 1000).toFixed(3)}
+              <Text color="gray.500"> / 1K tokens</Text>
+            </Text>
+          }
+        />
+        <SelectedModelLabeledInfo label="Speed" info={<Text>{stats.speed}</Text>} />
+      </SimpleGrid>
+    </VStack>
+  );
+};
+
+const SelectedModelLabeledInfo = ({
+  label,
+  info,
+  ...props
+}: {
+  label: string;
+  info: string | number | React.ReactElement;
+} & StackProps) => (
+  <GridItem>
+    <VStack alignItems="flex-start" {...props}>
+      <Text fontWeight="bold">{label}</Text>
+      <Text>{info}</Text>
+    </VStack>
+  </GridItem>
+);
--- a/src/components/SelectModelModal/SelectModelModal.tsx
+++ b/src/components/SelectModelModal/SelectModelModal.tsx
@@ -0,0 +1,77 @@
+import {
+  Button,
+  Modal,
+  ModalBody,
+  ModalCloseButton,
+  ModalContent,
+  ModalFooter,
+  ModalHeader,
+  ModalOverlay,
+  VStack,
+  Text,
+  Spinner,
+} from "@chakra-ui/react";
+import { useState } from "react";
+import { type SupportedModel } from "~/server/types";
+import { ModelStatsCard } from "./ModelStatsCard";
+import { SelectModelSearch } from "./SelectModelSearch";
+import { api } from "~/utils/api";
+import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
+
+export const SelectModelModal = ({
+  originalModel,
+  variantId,
+  onClose,
+}: {
+  originalModel: SupportedModel;
+  variantId: string;
+  onClose: () => void;
+}) => {
+  const [selectedModel, setSelectedModel] = useState<SupportedModel>(originalModel);
+  const utils = api.useContext();
+
+  const experiment = useExperiment();
+
+  const duplicateMutation = api.promptVariants.create.useMutation();
+
+  const [createNewVariant, creationInProgress] = useHandledAsyncCallback(async () => {
+    if (!experiment?.data?.id) return;
+    await duplicateMutation.mutateAsync({
+      experimentId: experiment?.data?.id,
+      variantId,
+      newModel: selectedModel,
+    });
+    await utils.promptVariants.list.invalidate();
+    onClose();
+  }, [duplicateMutation, experiment?.data?.id, variantId, onClose]);
+
+  return (
+    <Modal isOpen onClose={onClose} size={{ base: "xl", sm: "2xl", md: "3xl" }}>
+      <ModalOverlay />
+      <ModalContent w={1200}>
+        <ModalHeader>Select a New Model</ModalHeader>
+        <ModalCloseButton />
+        <ModalBody maxW="unset">
+          <VStack spacing={4}>
+            <ModelStatsCard label="ORIGINAL MODEL" model={originalModel} />
+            {originalModel !== selectedModel && (
+              <ModelStatsCard label="SELECTED MODEL" model={selectedModel} />
+            )}
+            <SelectModelSearch selectedModel={selectedModel} setSelectedModel={setSelectedModel} />
+          </VStack>
+        </ModalBody>
+
+        <ModalFooter>
+          <Button
+            colorScheme="blue"
+            onClick={createNewVariant}
+            w={20}
+            disabled={originalModel === selectedModel}
+          >
+            {creationInProgress ? <Spinner boxSize={4} /> : <Text>Continue</Text>}
+          </Button>
+        </ModalFooter>
+      </ModalContent>
+    </Modal>
+  );
+};
--- a/src/components/SelectModelModal/SelectModelSearch.tsx
+++ b/src/components/SelectModelModal/SelectModelSearch.tsx
@@ -0,0 +1,47 @@
+import { VStack, Text } from "@chakra-ui/react";
+import { type LegacyRef, useCallback } from "react";
+import Select, { type SingleValue } from "react-select";
+import { type SupportedModel } from "~/server/types";
+import { useElementDimensions } from "~/utils/hooks";
+
+const modelOptions: { value: SupportedModel; label: string }[] = [
+  { value: "gpt-3.5-turbo", label: "gpt-3.5-turbo" },
+  { value: "gpt-3.5-turbo-0613", label: "gpt-3.5-turbo-0613" },
+  { value: "gpt-3.5-turbo-16k", label: "gpt-3.5-turbo-16k" },
+  { value: "gpt-3.5-turbo-16k-0613", label: "gpt-3.5-turbo-16k-0613" },
+  { value: "gpt-4", label: "gpt-4" },
+  { value: "gpt-4-0613", label: "gpt-4-0613" },
+  { value: "gpt-4-32k", label: "gpt-4-32k" },
+  { value: "gpt-4-32k-0613", label: "gpt-4-32k-0613" },
+];
+
+export const SelectModelSearch = ({
+  selectedModel,
+  setSelectedModel,
+}: {
+  selectedModel: SupportedModel;
+  setSelectedModel: (model: SupportedModel) => void;
+}) => {
+  const handleSelection = useCallback(
+    (option: SingleValue<{ value: SupportedModel; label: string }>) => {
+      if (!option) return;
+      setSelectedModel(option.value);
+    },
+    [setSelectedModel],
+  );
+  const selectedOption = modelOptions.find((option) => option.value === selectedModel);
+
+  const [containerRef, containerDimensions] = useElementDimensions();
+
+  return (
+    <VStack ref={containerRef as LegacyRef<HTMLDivElement>} w="full">
+      <Text>Browse Models</Text>
+      <Select
+        styles={{ control: (provided) => ({ ...provided, width: containerDimensions?.width }) }}
+        value={selectedOption}
+        options={modelOptions}
+        onChange={handleSelection}
+      />
+    </VStack>
+  );
+};
--- a/src/components/VariantHeader/VariantHeader.tsx
+++ b/src/components/VariantHeader/VariantHeader.tsx
@@ -1,26 +1,13 @@
 import { useState, type DragEvent } from "react";
-import { type PromptVariant } from "./types";
+import { type PromptVariant } from "../OutputsTable/types";
 import { api } from "~/utils/api";
 import { useHandledAsyncCallback } from "~/utils/hooks";
-import {
-  Button,
-  HStack,
-  Icon,
-  Menu,
-  MenuButton,
-  MenuItem,
-  MenuList,
-  MenuDivider,
-  Text,
-  GridItem,
-  Spinner,
-} from "@chakra-ui/react"; // Changed here
-import { BsFillTrashFill, BsGear } from "react-icons/bs";
-import { FaRegClone } from "react-icons/fa";
-import { RiDraggable, RiExchangeFundsFill } from "react-icons/ri";
+import { HStack, Icon, GridItem } from "@chakra-ui/react"; // Changed here
+import { RiDraggable } from "react-icons/ri";
 import { cellPadding, headerMinHeight } from "../constants";
 import AutoResizeTextArea from "../AutoResizeTextArea";
-import { stickyHeaderStyle } from "./styles";
+import { stickyHeaderStyle } from "../OutputsTable/styles";
+import VariantHeaderMenuButton from "./VariantHeaderMenuButton";

 export default function VariantHeader(props: { variant: PromptVariant; canHide: boolean }) {
  const utils = api.useContext();
@@ -38,14 +25,6 @@ export default function VariantHeader(props: { variant: PromptVariant; canHide:
    }
  }, [updateMutation, props.variant.id, props.variant.label, label]);

-  const hideMutation = api.promptVariants.hide.useMutation();
-  const [onHide] = useHandledAsyncCallback(async () => {
-    await hideMutation.mutateAsync({
-      id: props.variant.id,
-    });
-    await utils.promptVariants.list.invalidate();
-  }, [hideMutation, props.variant.id]);
-
  const reorderMutation = api.promptVariants.reorder.useMutation();
  const [onReorder] = useHandledAsyncCallback(
    async (e: DragEvent<HTMLDivElement>) => {
@@ -64,21 +43,13 @@ export default function VariantHeader(props: { variant: PromptVariant; canHide:
  );

  const [menuOpen, setMenuOpen] = useState(false);
-  const duplicateMutation = api.promptVariants.create.useMutation();
-
-  const [duplicateVariant, duplicationInProgress] = useHandledAsyncCallback(async () => {
-    await duplicateMutation.mutateAsync({
-      experimentId: props.variant.experimentId,
-      variantId: props.variant.id,
-    });
-    await utils.promptVariants.list.invalidate();
-  }, [duplicateMutation, props.variant.experimentId, props.variant.id]);

  return (
    <GridItem
      padding={0}
      sx={{
        ...stickyHeaderStyle,
+        // Ensure that the menu always appears above the sticky header of other variants
        zIndex: menuOpen ? "dropdown" : stickyHeaderStyle.zIndex,
      }}
      borderTopWidth={1}
@@ -129,42 +100,12 @@ export default function VariantHeader(props: { variant: PromptVariant; canHide:
          onMouseEnter={() => setIsInputHovered(true)}
          onMouseLeave={() => setIsInputHovered(false)}
        />
-
-        <Menu
-          z-index="dropdown"
-          onOpen={() => setMenuOpen(true)}
-          onClose={() => setMenuOpen(false)}
-        >
-          {duplicationInProgress ? (
-            <Spinner boxSize={4} mx={3} my={3} />
-          ) : (
-            <MenuButton>
-              <Button variant="ghost">
-                <Icon as={BsGear} />
-              </Button>
-            </MenuButton>
-          )}
-
-          <MenuList mt={-3} fontSize="md">
-            <MenuItem icon={<Icon as={FaRegClone} boxSize={4} w={5} />} onClick={duplicateVariant}>
-              Duplicate
-            </MenuItem>
-            <MenuItem icon={<Icon as={RiExchangeFundsFill} boxSize={5} />}>Change Model</MenuItem>
-            {props.canHide && (
-              <>
-                <MenuDivider />
-                <MenuItem
-                  onClick={onHide}
-                  icon={<Icon as={BsFillTrashFill} boxSize={5} />}
-                  color="red.600"
-                  _hover={{ backgroundColor: "red.50" }}
-                >
-                  <Text>Hide</Text>
-                </MenuItem>
-              </>
-            )}
-          </MenuList>
-        </Menu>
+        <VariantHeaderMenuButton
+          variant={props.variant}
+          canHide={props.canHide}
+          menuOpen={menuOpen}
+          setMenuOpen={setMenuOpen}
+        />
      </HStack>
    </GridItem>
  );
--- a/src/components/VariantHeader/VariantHeaderMenuButton.tsx
+++ b/src/components/VariantHeader/VariantHeaderMenuButton.tsx
@@ -0,0 +1,102 @@
+import { type PromptVariant } from "../OutputsTable/types";
+import { api } from "~/utils/api";
+import { useHandledAsyncCallback } from "~/utils/hooks";
+import {
+  Button,
+  Icon,
+  Menu,
+  MenuButton,
+  MenuItem,
+  MenuList,
+  MenuDivider,
+  Text,
+  Spinner,
+} from "@chakra-ui/react"; // Changed here
+import { BsFillTrashFill, BsGear } from "react-icons/bs";
+import { FaRegClone } from "react-icons/fa";
+import { RiExchangeFundsFill } from "react-icons/ri";
+import { useState } from "react";
+import { SelectModelModal } from "../SelectModelModal/SelectModelModal";
+import { type SupportedModel } from "~/server/types";
+
+export default function VariantHeaderMenuButton({
+  variant,
+  canHide,
+  menuOpen,
+  setMenuOpen,
+}: {
+  variant: PromptVariant;
+  canHide: boolean;
+  menuOpen: boolean;
+  setMenuOpen: (open: boolean) => void;
+}) {
+  const utils = api.useContext();
+
+  const duplicateMutation = api.promptVariants.create.useMutation();
+
+  const [duplicateVariant, duplicationInProgress] = useHandledAsyncCallback(async () => {
+    await duplicateMutation.mutateAsync({
+      experimentId: variant.experimentId,
+      variantId: variant.id,
+    });
+    await utils.promptVariants.list.invalidate();
+  }, [duplicateMutation, variant.experimentId, variant.id]);
+
+  const hideMutation = api.promptVariants.hide.useMutation();
+  const [onHide] = useHandledAsyncCallback(async () => {
+    await hideMutation.mutateAsync({
+      id: variant.id,
+    });
+    await utils.promptVariants.list.invalidate();
+  }, [hideMutation, variant.id]);
+
+  const [selectModelModalOpen, setSelectModelModalOpen] = useState(false);
+
+  return (
+    <>
+      <Menu isOpen={menuOpen} onOpen={() => setMenuOpen(true)} onClose={() => setMenuOpen(false)}>
+        {duplicationInProgress ? (
+          <Spinner boxSize={4} mx={3} my={3} />
+        ) : (
+          <MenuButton>
+            <Button variant="ghost">
+              <Icon as={BsGear} />
+            </Button>
+          </MenuButton>
+        )}
+
+        <MenuList mt={-3} fontSize="md">
+          <MenuItem icon={<Icon as={FaRegClone} boxSize={4} w={5} />} onClick={duplicateVariant}>
+            Duplicate
+          </MenuItem>
+          <MenuItem
+            icon={<Icon as={RiExchangeFundsFill} boxSize={5} />}
+            onClick={() => setSelectModelModalOpen(true)}
+          >
+            Change Model
+          </MenuItem>
+          {canHide && (
+            <>
+              <MenuDivider />
+              <MenuItem
+                onClick={onHide}
+                icon={<Icon as={BsFillTrashFill} boxSize={5} />}
+                color="red.600"
+                _hover={{ backgroundColor: "red.50" }}
+              >
+                <Text>Hide</Text>
+              </MenuItem>
+            </>
+          )}
+        </MenuList>
+      </Menu>
+      {selectModelModalOpen && (
+        <SelectModelModal
+          originalModel={variant.model as SupportedModel}
+          variantId={variant.id}
+          onClose={() => setSelectModelModalOpen(false)}
+        />
+      )}
+    </>
+  );
+}