Fullscreen editor

* Add support for switching to Llama models

* Fix prettier

.env.example

@@ -17,6 +17,9 @@ DATABASE_URL="postgresql://postgres:postgres@localhost:5432/openpipe?schema=publ
 # https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key
 OPENAI_API_KEY=""
 
+# Replicate API token. Create a token here: https://replicate.com/account/api-tokens
+REPLICATE_API_TOKEN=""
+
 NEXT_PUBLIC_SOCKET_URL="http://localhost:3318"
 
 # Next Auth

README.md

@@ -43,7 +43,8 @@ Natively supports [OpenAI function calls](https://openai.com/blog/function-calli
 
 ## Supported Models
 
-OpenPipe currently supports GPT-3.5 and GPT-4. Wider model support is planned.
+- All models available through the OpenAI [chat completion API](https://platform.openai.com/docs/guides/gpt/chat-completions-api)
+- Llama2 [7b chat](https://replicate.com/a16z-infra/llama7b-v2-chat), [13b chat](https://replicate.com/a16z-infra/llama13b-v2-chat), [70b chat](https://replicate.com/replicate/llama70b-v2-chat).
 
 ## Running Locally
 

src/components/AutoResizeTextArea.tsx

@@ -1,19 +1,22 @@
 import { Textarea, type TextareaProps } from "@chakra-ui/react";
 import ResizeTextarea from "react-textarea-autosize";
-import React from "react";
+import React, { useLayoutEffect, useState } from "react";
 
 export const AutoResizeTextarea: React.ForwardRefRenderFunction<
   HTMLTextAreaElement,
   TextareaProps & { minRows?: number }
-> = (props, ref) => {
+> = ({ minRows = 1, overflowY = "hidden", ...props }, ref) => {
+  const [isRerendered, setIsRerendered] = useState(false);
+  useLayoutEffect(() => setIsRerendered(true), []);
+
   return (
     <Textarea
       minH="unset"
-      overflow="hidden"
+      minRows={minRows}
+      overflowY={isRerendered ? overflowY : "hidden"}
       w="100%"
       resize="none"
       ref={ref}
-      minRows={1}
       transition="height none"
       as={ResizeTextarea}
       {...props}

src/components/ChangeModelModal/ChangeModelModal.tsx

@@ -0,0 +1,135 @@
+import {
+  Button,
+  Modal,
+  ModalBody,
+  ModalCloseButton,
+  ModalContent,
+  ModalFooter,
+  ModalHeader,
+  ModalOverlay,
+  VStack,
+  Text,
+  Spinner,
+  HStack,
+  Icon,
+} from "@chakra-ui/react";
+import { RiExchangeFundsFill } from "react-icons/ri";
+import { useState } from "react";
+import { ModelStatsCard } from "./ModelStatsCard";
+import { ModelSearch } from "./ModelSearch";
+import { api } from "~/utils/api";
+import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
+import CompareFunctions from "../RefinePromptModal/CompareFunctions";
+import { type PromptVariant } from "@prisma/client";
+import { isObject, isString } from "lodash-es";
+import { type Model, type SupportedProvider } from "~/modelProviders/types";
+import frontendModelProviders from "~/modelProviders/frontendModelProviders";
+import { keyForModel } from "~/utils/utils";
+
+export const ChangeModelModal = ({
+  variant,
+  onClose,
+}: {
+  variant: PromptVariant;
+  onClose: () => void;
+}) => {
+  const originalModelProviderName = variant.modelProvider as SupportedProvider;
+  const originalModelProvider = frontendModelProviders[originalModelProviderName];
+  const originalModel = originalModelProvider.models[variant.model] as Model;
+  const [selectedModel, setSelectedModel] = useState<Model>(originalModel);
+  const [convertedModel, setConvertedModel] = useState<Model | undefined>(undefined);
+  const utils = api.useContext();
+
+  const experiment = useExperiment();
+
+  const { mutateAsync: getModifiedPromptMutateAsync, data: modifiedPromptFn } =
+    api.promptVariants.getModifiedPromptFn.useMutation();
+
+  const [getModifiedPromptFn, modificationInProgress] = useHandledAsyncCallback(async () => {
+    if (!experiment) return;
+
+    await getModifiedPromptMutateAsync({
+      id: variant.id,
+      newModel: selectedModel,
+    });
+    setConvertedModel(selectedModel);
+  }, [getModifiedPromptMutateAsync, onClose, experiment, variant, selectedModel]);
+
+  const replaceVariantMutation = api.promptVariants.replaceVariant.useMutation();
+
+  const [replaceVariant, replacementInProgress] = useHandledAsyncCallback(async () => {
+    if (
+      !variant.experimentId ||
+      !modifiedPromptFn ||
+      (isObject(modifiedPromptFn) && "status" in modifiedPromptFn)
+    )
+      return;
+    await replaceVariantMutation.mutateAsync({
+      id: variant.id,
+      constructFn: modifiedPromptFn,
+    });
+    await utils.promptVariants.list.invalidate();
+    onClose();
+  }, [replaceVariantMutation, variant, onClose, modifiedPromptFn]);
+
+  const originalModelLabel = keyForModel(originalModel);
+  const selectedModelLabel = keyForModel(selectedModel);
+  const convertedModelLabel = convertedModel ? keyForModel(convertedModel) : undefined;
+
+  return (
+    <Modal
+      isOpen
+      onClose={onClose}
+      size={{ base: "xl", sm: "2xl", md: "3xl", lg: "5xl", xl: "7xl" }}
+    >
+      <ModalOverlay />
+      <ModalContent w={1200}>
+        <ModalHeader>
+          <HStack>
+            <Icon as={RiExchangeFundsFill} />
+            <Text>Change Model</Text>
+          </HStack>
+        </ModalHeader>
+        <ModalCloseButton />
+        <ModalBody maxW="unset">
+          <VStack spacing={8}>
+            <ModelStatsCard label="Original Model" model={originalModel} />
+            {originalModelLabel !== selectedModelLabel && (
+              <ModelStatsCard label="New Model" model={selectedModel} />
+            )}
+            <ModelSearch selectedModel={selectedModel} setSelectedModel={setSelectedModel} />
+            {isString(modifiedPromptFn) && (
+              <CompareFunctions
+                originalFunction={variant.constructFn}
+                newFunction={modifiedPromptFn}
+                leftTitle={originalModelLabel}
+                rightTitle={convertedModelLabel}
+              />
+            )}
+          </VStack>
+        </ModalBody>
+
+        <ModalFooter>
+          <HStack>
+            <Button
+              colorScheme="gray"
+              onClick={getModifiedPromptFn}
+              minW={24}
+              isDisabled={originalModel === selectedModel || modificationInProgress}
+            >
+              {modificationInProgress ? <Spinner boxSize={4} /> : <Text>Convert</Text>}
+            </Button>
+            <Button
+              colorScheme="blue"
+              onClick={replaceVariant}
+              minW={24}
+              isDisabled={!convertedModel || modificationInProgress || replacementInProgress}
+            >
+              {replacementInProgress ? <Spinner boxSize={4} /> : <Text>Accept</Text>}
+            </Button>
+          </HStack>
+        </ModalFooter>
+      </ModalContent>
+    </Modal>
+  );
+};

src/components/ChangeModelModal/ModelSearch.tsx

@@ -0,0 +1,50 @@
+import { VStack, Text } from "@chakra-ui/react";
+import { type LegacyRef, useCallback } from "react";
+import Select, { type SingleValue } from "react-select";
+import { useElementDimensions } from "~/utils/hooks";
+
+import frontendModelProviders from "~/modelProviders/frontendModelProviders";
+import { type Model } from "~/modelProviders/types";
+import { keyForModel } from "~/utils/utils";
+
+const modelOptions: { label: string; value: Model }[] = [];
+
+for (const [_, providerValue] of Object.entries(frontendModelProviders)) {
+  for (const [_, modelValue] of Object.entries(providerValue.models)) {
+    modelOptions.push({
+      label: keyForModel(modelValue),
+      value: modelValue,
+    });
+  }
+}
+
+export const ModelSearch = ({
+  selectedModel,
+  setSelectedModel,
+}: {
+  selectedModel: Model;
+  setSelectedModel: (model: Model) => void;
+}) => {
+  const handleSelection = useCallback(
+    (option: SingleValue<{ label: string; value: Model }>) => {
+      if (!option) return;
+      setSelectedModel(option.value);
+    },
+    [setSelectedModel],
+  );
+  const selectedOption = modelOptions.find((option) => option.label === keyForModel(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>
+  );
+};

src/components/ChangeModelModal/ModelStatsCard.tsx

@@ -7,11 +7,9 @@ import {
   SimpleGrid,
   Link,
 } from "@chakra-ui/react";
-import { modelStats } from "~/modelProviders/modelStats";
-import { type SupportedModel } from "~/server/types";
+import { type Model } from "~/modelProviders/types";
 
-export const ModelStatsCard = ({ label, model }: { label: string; model: SupportedModel }) => {
-  const stats = modelStats[model];
+export const ModelStatsCard = ({ label, model }: { label: string; model: Model }) => {
   return (
     <VStack w="full" align="start">
       <Text fontWeight="bold" fontSize="sm" textTransform="uppercase">
@@ -22,14 +20,14 @@ export const ModelStatsCard = ({ label, model }: { label: string; model: Support
         <HStack w="full" align="flex-start">
           <Text flex={1} fontSize="lg">
             <Text as="span" color="gray.600">
-              {stats.provider} /{" "}
+              {model.provider} /{" "}
             </Text>
             <Text as="span" fontWeight="bold" color="gray.900">
-              {model}
+              {model.name}
             </Text>
           </Text>
           <Link
-            href={stats.learnMoreUrl}
+            href={model.learnMoreUrl}
             isExternal
             color="blue.500"
             fontWeight="bold"
@@ -46,26 +44,41 @@ export const ModelStatsCard = ({ label, model }: { label: string; model: Support
           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>} />
+          <SelectedModelLabeledInfo label="Context Window" info={model.contextWindow} />
+          {model.promptTokenPrice && (
+            <SelectedModelLabeledInfo
+              label="Input"
+              info={
+                <Text>
+                  ${(model.promptTokenPrice * 1000).toFixed(3)}
+                  <Text color="gray.500"> / 1K tokens</Text>
+                </Text>
+              }
+            />
+          )}
+          {model.completionTokenPrice && (
+            <SelectedModelLabeledInfo
+              label="Output"
+              info={
+                <Text>
+                  ${(model.completionTokenPrice * 1000).toFixed(3)}
+                  <Text color="gray.500"> / 1K tokens</Text>
+                </Text>
+              }
+            />
+          )}
+          {model.pricePerSecond && (
+            <SelectedModelLabeledInfo
+              label="Price"
+              info={
+                <Text>
+                  ${model.pricePerSecond.toFixed(3)}
+                  <Text color="gray.500"> / second</Text>
+                </Text>
+              }
+            />
+          )}
+          <SelectedModelLabeledInfo label="Speed" info={<Text>{model.speed}</Text>} />
         </SimpleGrid>
       </VStack>
     </VStack>

src/components/OutputsTable/FloatingLabelInput.tsx

@@ -18,11 +18,9 @@ export const FloatingLabelInput = ({
         transform={isFocused || !!value ? "translateY(-50%)" : "translateY(0)"}
         fontSize={isFocused || !!value ? "12px" : "16px"}
         transition="all 0.15s"
-        zIndex="100"
+        zIndex="5"
         bg="white"
         px={1}
-        mt={0}
-        mb={2}
         lineHeight="1"
         pointerEvents="none"
         color={isFocused ? "blue.500" : "gray.500"}

src/components/OutputsTable/NewScenarioButton.tsx

@@ -49,7 +49,11 @@ export default function NewScenarioButton() {
         Add Scenario
       </StyledButton>
       <StyledButton onClick={onAutogenerate}>
-        <Icon as={autogenerating ? Spinner : BsPlus} boxSize={6} mr={autogenerating ? 1 : 0} />
+        <Icon
+          as={autogenerating ? Spinner : BsPlus}
+          boxSize={autogenerating ? 4 : 6}
+          mr={autogenerating ? 2 : 0}
+        />
         Autogenerate Scenario
       </StyledButton>
     </HStack>

src/components/OutputsTable/OutputCell/OutputCell.tsx

@@ -10,7 +10,7 @@ import useSocket from "~/utils/useSocket";
 import { OutputStats } from "./OutputStats";
 import { ErrorHandler } from "./ErrorHandler";
 import { CellOptions } from "./CellOptions";
-import modelProvidersFrontend from "~/modelProviders/modelProvidersFrontend";
+import frontendModelProviders from "~/modelProviders/frontendModelProviders";
 
 export default function OutputCell({
   scenario,
@@ -40,7 +40,7 @@ export default function OutputCell({
   );
 
   const provider =
-    modelProvidersFrontend[variant.modelProvider as keyof typeof modelProvidersFrontend];
+    frontendModelProviders[variant.modelProvider as keyof typeof frontendModelProviders];
 
   type OutputSchema = Parameters<typeof provider.normalizeOutput>[0];
 
@@ -88,11 +88,9 @@ export default function OutputCell({
   }
 
   const normalizedOutput = modelOutput
-    ? // @ts-expect-error TODO FIX ASAP
-      provider.normalizeOutput(modelOutput.output as unknown as OutputSchema)
+    ? provider.normalizeOutput(modelOutput.output)
     : streamedMessage
-    ? // @ts-expect-error TODO FIX ASAP
-      provider.normalizeOutput(streamedMessage)
+    ? provider.normalizeOutput(streamedMessage)
     : null;
 
   if (modelOutput && normalizedOutput?.type === "json") {

src/components/OutputsTable/VariantEditor.tsx

@@ -1,17 +1,47 @@
-import { Box, Button, HStack, Spinner, Tooltip, useToast, Text } from "@chakra-ui/react";
+import {
+  Box,
+  Button,
+  HStack,
+  Spinner,
+  Tooltip,
+  useToast,
+  Text,
+  IconButton,
+} from "@chakra-ui/react";
 import { useRef, useEffect, useState, useCallback } from "react";
 import { useExperimentAccess, useHandledAsyncCallback, useModifierKeyLabel } from "~/utils/hooks";
 import { type PromptVariant } from "./types";
 import { api } from "~/utils/api";
 import { useAppStore } from "~/state/store";
+import { FiMaximize, FiMinimize } from "react-icons/fi";
+import { editorBackground } from "~/state/sharedVariantEditor.slice";
 
 export default function VariantEditor(props: { variant: PromptVariant }) {
   const { canModify } = useExperimentAccess();
   const monaco = useAppStore.use.sharedVariantEditor.monaco();
   const editorRef = useRef<ReturnType<NonNullable<typeof monaco>["editor"]["create"]> | null>(null);
+  const containerRef = useRef<HTMLDivElement | null>(null);
   const [editorId] = useState(() => `editor_${Math.random().toString(36).substring(7)}`);
   const [isChanged, setIsChanged] = useState(false);
 
+  const [isFullscreen, setIsFullscreen] = useState(false);
+
+  const toggleFullscreen = useCallback(() => {
+    setIsFullscreen((prev) => !prev);
+    editorRef.current?.focus();
+  }, [setIsFullscreen]);
+
+  useEffect(() => {
+    const handleEsc = (event: KeyboardEvent) => {
+      if (event.key === "Escape" && isFullscreen) {
+        toggleFullscreen();
+      }
+    };
+
+    window.addEventListener("keydown", handleEsc);
+    return () => window.removeEventListener("keydown", handleEsc);
+  }, [isFullscreen, toggleFullscreen]);
+
   const lastSavedFn = props.variant.constructFn;
 
   const modifierKey = useModifierKeyLabel();
@@ -99,11 +129,23 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
         readOnly: !canModify,
       });
 
+      // Workaround because otherwise the commands only work on whatever
+      // editor was loaded on the page last.
+      // https://github.com/microsoft/monaco-editor/issues/2947#issuecomment-1422265201
       editorRef.current.onDidFocusEditorText(() => {
-        // Workaround because otherwise the command only works on whatever
-        // editor was loaded on the page last.
-        // https://github.com/microsoft/monaco-editor/issues/2947#issuecomment-1422265201
-        editorRef.current?.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.Enter, onSave);
+        editorRef.current?.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyS, onSave);
+
+        editorRef.current?.addCommand(
+          monaco.KeyMod.CtrlCmd | monaco.KeyMod.Shift | monaco.KeyCode.KeyF,
+          toggleFullscreen,
+        );
+
+        // Exit fullscreen with escape
+        editorRef.current?.addCommand(monaco.KeyCode.Escape, () => {
+          if (isFullscreen) {
+            toggleFullscreen();
+          }
+        });
       });
 
       editorRef.current.onDidChangeModelContent(checkForChanges);
@@ -132,8 +174,40 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
   }, [canModify]);
 
   return (
-    <Box w="100%" pos="relative">
-      <div id={editorId} style={{ height: "400px", width: "100%" }}></div>
+    <Box
+      w="100%"
+      ref={containerRef}
+      sx={
+        isFullscreen
+          ? {
+              position: "fixed",
+              top: 0,
+              left: 0,
+              right: 0,
+              bottom: 0,
+            }
+          : { h: "400px", w: "100%" }
+      }
+      bgColor={editorBackground}
+      zIndex={isFullscreen ? 1000 : "unset"}
+      pos="relative"
+      _hover={{ ".fullscreen-toggle": { opacity: 1 } }}
+    >
+      <Box id={editorId} w="100%" h="100%" />
+      <Tooltip label={`${modifierKey} + ⇧ + F`}>
+        <IconButton
+          className="fullscreen-toggle"
+          aria-label="Minimize"
+          icon={isFullscreen ? <FiMinimize /> : <FiMaximize />}
+          position="absolute"
+          top={2}
+          right={2}
+          onClick={toggleFullscreen}
+          opacity={0}
+          transition="opacity 0.2s"
+        />
+      </Tooltip>
+
       {isChanged && (
         <HStack pos="absolute" bottom={2} right={2}>
           <Button
@@ -146,7 +220,7 @@ export default function VariantEditor(props: { variant: PromptVariant }) {
           >
             Reset
           </Button>
-          <Tooltip label={`${modifierKey} + Enter`}>
+          <Tooltip label={`${modifierKey} + S`}>
             <Button size="sm" onClick={onSave} colorScheme="blue" w={16} disabled={saveInProgress}>
               {saveInProgress ? <Spinner boxSize={4} /> : <Text>Save</Text>}
             </Button>

src/components/OutputsTable/styles.ts

@@ -4,5 +4,5 @@ export const stickyHeaderStyle: SystemStyleObject = {
   position: "sticky",
   top: "0",
   backgroundColor: "#fff",
-  zIndex: 1,
+  zIndex: 10,
 };

src/components/RefinePromptModal/CompareFunctions.tsx

@@ -1,4 +1,4 @@
-import { HStack, VStack, useBreakpointValue } from "@chakra-ui/react";
+import { type StackProps, VStack, useBreakpointValue } from "@chakra-ui/react";
 import React from "react";
 import DiffViewer, { DiffMethod } from "react-diff-viewer";
 import Prism from "prismjs";
@@ -19,10 +19,15 @@ const highlightSyntax = (str: string) => {
 const CompareFunctions = ({
   originalFunction,
   newFunction = "",
+  leftTitle = "Original",
+  rightTitle = "Modified",
+  ...props
 }: {
   originalFunction: string;
   newFunction?: string;
-}) => {
+  leftTitle?: string;
+  rightTitle?: string;
+} & StackProps) => {
   const showSplitView = useBreakpointValue(
     {
       base: false,
@@ -34,22 +39,20 @@ const CompareFunctions = ({
   );
 
   return (
-    <HStack w="full" spacing={5}>
-      <VStack w="full" spacing={4} maxH="40vh" fontSize={12} lineHeight={1} overflowY="auto">
-        <DiffViewer
-          oldValue={originalFunction}
-          newValue={newFunction || originalFunction}
-          splitView={showSplitView}
-          hideLineNumbers={!showSplitView}
-          leftTitle="Original"
-          rightTitle={newFunction ? "Modified" : "Unmodified"}
-          disableWordDiff={true}
-          compareMethod={DiffMethod.CHARS}
-          renderContent={highlightSyntax}
-          showDiffOnly={false}
-        />
-      </VStack>
-    </HStack>
+    <VStack w="full" spacing={4} fontSize={12} lineHeight={1} overflowY="auto" {...props}>
+      <DiffViewer
+        oldValue={originalFunction}
+        newValue={newFunction || originalFunction}
+        splitView={showSplitView}
+        hideLineNumbers={!showSplitView}
+        leftTitle={leftTitle}
+        rightTitle={rightTitle}
+        disableWordDiff={true}
+        compareMethod={DiffMethod.CHARS}
+        renderContent={highlightSyntax}
+        showDiffOnly={false}
+      />
+    </VStack>
   );
 };
 

src/components/RefinePromptModal/CustomInstructionsInput.tsx

@@ -56,7 +56,6 @@ export const CustomInstructionsInput = ({
           minW="unset"
           size="sm"
           onClick={() => onSubmit()}
-          disabled={!instructions}
           variant={instructions ? "solid" : "ghost"}
           mr={4}
           borderRadius="8"

src/components/RefinePromptModal/RefineOption.tsx

@@ -1,22 +1,22 @@
 import { HStack, Icon, Heading, Text, VStack, GridItem } from "@chakra-ui/react";
 import { type IconType } from "react-icons";
-import { refineOptions, type RefineOptionLabel } from "./refineOptions";
 
 export const RefineOption = ({
   label,
-  activeLabel,
   icon,
+  desciption,
+  activeLabel,
   onClick,
   loading,
 }: {
-  label: RefineOptionLabel;
-  activeLabel: RefineOptionLabel | undefined;
+  label: string;
   icon: IconType;
-  onClick: (label: RefineOptionLabel) => void;
+  desciption: string;
+  activeLabel: string | undefined;
+  onClick: (label: string) => void;
   loading: boolean;
 }) => {
   const isActive = activeLabel === label;
-  const desciption = refineOptions[label].description;
 
   return (
     <GridItem w="80" h="44">

src/components/RefinePromptModal/RefinePromptModal.tsx

@@ -15,17 +15,16 @@ import {
   SimpleGrid,
 } from "@chakra-ui/react";
 import { BsStars } from "react-icons/bs";
-import { VscJson } from "react-icons/vsc";
-import { TfiThought } from "react-icons/tfi";
 import { api } from "~/utils/api";
 import { useHandledAsyncCallback } from "~/utils/hooks";
 import { type PromptVariant } from "@prisma/client";
 import { useState } from "react";
 import CompareFunctions from "./CompareFunctions";
 import { CustomInstructionsInput } from "./CustomInstructionsInput";
-import { type RefineOptionLabel, refineOptions } from "./refineOptions";
+import { type RefineOptionInfo, refineOptions } from "./refineOptions";
 import { RefineOption } from "./RefineOption";
 import { isObject, isString } from "lodash-es";
+import { type SupportedProvider } from "~/modelProviders/types";
 
 export const RefinePromptModal = ({
   variant,
@@ -36,25 +35,29 @@ export const RefinePromptModal = ({
 }) => {
   const utils = api.useContext();
 
-  const { mutateAsync: getRefinedPromptMutateAsync, data: refinedPromptFn } =
-    api.promptVariants.getRefinedPromptFn.useMutation();
+  const providerRefineOptions = refineOptions[variant.modelProvider as SupportedProvider];
+
+  const { mutateAsync: getModifiedPromptMutateAsync, data: refinedPromptFn } =
+    api.promptVariants.getModifiedPromptFn.useMutation();
   const [instructions, setInstructions] = useState<string>("");
 
-  const [activeRefineOptionLabel, setActiveRefineOptionLabel] = useState<
-    RefineOptionLabel | undefined
-  >(undefined);
+  const [activeRefineOptionLabel, setActiveRefineOptionLabel] = useState<string | undefined>(
+    undefined,
+  );
 
-  const [getRefinedPromptFn, refiningInProgress] = useHandledAsyncCallback(
-    async (label?: RefineOptionLabel) => {
+  const [getModifiedPromptFn, modificationInProgress] = useHandledAsyncCallback(
+    async (label?: string) => {
       if (!variant.experimentId) return;
-      const updatedInstructions = label ? refineOptions[label].instructions : instructions;
+      const updatedInstructions = label
+        ? (providerRefineOptions[label] as RefineOptionInfo).instructions
+        : instructions;
       setActiveRefineOptionLabel(label);
-      await getRefinedPromptMutateAsync({
+      await getModifiedPromptMutateAsync({
         id: variant.id,
         instructions: updatedInstructions,
       });
     },
-    [getRefinedPromptMutateAsync, onClose, variant, instructions, setActiveRefineOptionLabel],
+    [getModifiedPromptMutateAsync, onClose, variant, instructions, setActiveRefineOptionLabel],
   );
 
   const replaceVariantMutation = api.promptVariants.replaceVariant.useMutation();
@@ -75,7 +78,11 @@ export const RefinePromptModal = ({
   }, [replaceVariantMutation, variant, onClose, refinedPromptFn]);
 
   return (
-    <Modal isOpen onClose={onClose} size={{ base: "xl", sm: "2xl", md: "7xl" }}>
+    <Modal
+      isOpen
+      onClose={onClose}
+      size={{ base: "xl", sm: "2xl", md: "3xl", lg: "5xl", xl: "7xl" }}
+    >
       <ModalOverlay />
       <ModalContent w={1200}>
         <ModalHeader>
@@ -88,35 +95,37 @@ export const RefinePromptModal = ({
         <ModalBody maxW="unset">
           <VStack spacing={8}>
             <VStack spacing={4}>
-              <SimpleGrid columns={{ base: 1, md: 2 }} spacing={8}>
-                <RefineOption
-                  label="Convert to function call"
-                  activeLabel={activeRefineOptionLabel}
-                  icon={VscJson}
-                  onClick={getRefinedPromptFn}
-                  loading={refiningInProgress}
-                />
-                <RefineOption
-                  label="Add chain of thought"
-                  activeLabel={activeRefineOptionLabel}
-                  icon={TfiThought}
-                  onClick={getRefinedPromptFn}
-                  loading={refiningInProgress}
-                />
-              </SimpleGrid>
-              <HStack>
-                <Text color="gray.500">or</Text>
-              </HStack>
+              {Object.keys(providerRefineOptions).length && (
+                <>
+                  <SimpleGrid columns={{ base: 1, md: 2 }} spacing={8}>
+                    {Object.keys(providerRefineOptions).map((label) => (
+                      <RefineOption
+                        key={label}
+                        label={label}
+                        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+                        icon={providerRefineOptions[label]!.icon}
+                        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+                        desciption={providerRefineOptions[label]!.description}
+                        activeLabel={activeRefineOptionLabel}
+                        onClick={getModifiedPromptFn}
+                        loading={modificationInProgress}
+                      />
+                    ))}
+                  </SimpleGrid>
+                  <Text color="gray.500">or</Text>
+                </>
+              )}
               <CustomInstructionsInput
                 instructions={instructions}
                 setInstructions={setInstructions}
-                loading={refiningInProgress}
-                onSubmit={getRefinedPromptFn}
+                loading={modificationInProgress}
+                onSubmit={getModifiedPromptFn}
               />
             </VStack>
             <CompareFunctions
               originalFunction={variant.constructFn}
               newFunction={isString(refinedPromptFn) ? refinedPromptFn : undefined}
+              maxH="40vh"
             />
           </VStack>
         </ModalBody>
@@ -124,12 +133,10 @@ export const RefinePromptModal = ({
         <ModalFooter>
           <HStack spacing={4}>
             <Button
+              colorScheme="blue"
               onClick={replaceVariant}
               minW={24}
-              disabled={replacementInProgress || !refinedPromptFn}
-              _disabled={{
-                bgColor: "blue.500",
-              }}
+              isDisabled={replacementInProgress || !refinedPromptFn}
             >
               {replacementInProgress ? <Spinner boxSize={4} /> : <Text>Accept</Text>}
             </Button>

src/components/RefinePromptModal/refineOptions.ts

@@ -1,18 +1,22 @@
 // Super hacky, but we'll redo the organization when we have more models
 
-export type RefineOptionLabel = "Add chain of thought" | "Convert to function call";
+import { type SupportedProvider } from "~/modelProviders/types";
+import { VscJson } from "react-icons/vsc";
+import { TfiThought } from "react-icons/tfi";
+import { type IconType } from "react-icons";
+
+export type RefineOptionInfo = { icon: IconType; description: string; instructions: string };
+
+export const refineOptions: Record<SupportedProvider, { [key: string]: RefineOptionInfo }> = {
+  "openai/ChatCompletion": {
+    "Add chain of thought": {
+      icon: VscJson,
+      description: "Asking the model to plan its answer can increase accuracy.",
+      instructions: `Adding chain of thought means asking the model to think about its answer before it gives it to you. This is useful for getting more accurate answers. Do not add an assistant message.
 
-export const refineOptions: Record<
-  RefineOptionLabel,
-  { description: string; instructions: string }
-> = {
-  "Add chain of thought": {
-    description: "Asking the model to plan its answer can increase accuracy.",
-    instructions: `Adding chain of thought means asking the model to think about its answer before it gives it to you. This is useful for getting more accurate answers. Do not add an assistant message.
-      
     This is what a prompt looks like before adding chain of thought:
-    
-    prompt = {
+
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-4",
         stream: true,
         messages: [
@@ -25,11 +29,11 @@ export const refineOptions: Record<
             content: \`This is the user's message: \${scenario.user_message}. Return "positive" or "negative" or "neutral"\`,
             },
         ],
-    };
+    });
 
     This is what one looks like after adding chain of thought:
 
-    prompt = {
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-4",
         stream: true,
         messages: [
@@ -42,22 +46,22 @@ export const refineOptions: Record<
             content: \`This is the user's message: \${scenario.user_message}. Return "positive" or "negative" or "neutral". Explain your answer before you give a score, then return the score on a new line.\`,
             },
         ],
-    };
+    });
 
     Here's another example:
 
     Before:
 
-    prompt = {
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-3.5-turbo",
         messages: [
           {
             role: "user",
             content: \`Title: \${scenario.title}
       Body: \${scenario.body}
-      
+
       Need: \${scenario.need}
-      
+
       Rate likelihood on 1-3 scale.\`,
           },
         ],
@@ -78,20 +82,20 @@ export const refineOptions: Record<
         function_call: {
           name: "score_post",
         },
-      };
+      });
 
     After:
 
-    prompt = {
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-3.5-turbo",
         messages: [
           {
             role: "user",
             content: \`Title: \${scenario.title}
       Body: \${scenario.body}
-      
+
       Need: \${scenario.need}
-      
+
       Rate likelihood on 1-3 scale. Provide an explanation, but always provide a score afterward.\`,
           },
         ],
@@ -115,17 +119,18 @@ export const refineOptions: Record<
         function_call: {
           name: "score_post",
         },
-      };
+      });
 
     Add chain of thought to the original prompt.`,
-  },
-  "Convert to function call": {
-    description: "Use function calls to get output from the model in a more structured way.",
-    instructions: `OpenAI functions are a specialized way for an LLM to return output.
-    
+    },
+    "Convert to function call": {
+      icon: TfiThought,
+      description: "Use function calls to get output from the model in a more structured way.",
+      instructions: `OpenAI functions are a specialized way for an LLM to return output.
+
     This is what a prompt looks like before adding a function:
-    
-    prompt = {
+
+    definePrompt("openai/ChatCompletion", {
       model: "gpt-4",
       stream: true,
       messages: [
@@ -138,11 +143,11 @@ export const refineOptions: Record<
           content: \`This is the user's message: \${scenario.user_message}. Return "positive" or "negative" or "neutral"\`,
         },
       ],
-    };
-  
+    });
+
     This is what one looks like after adding a function:
-  
-    prompt = {
+
+    definePrompt("openai/ChatCompletion", {
       model: "gpt-4",
       stream: true,
       messages: [
@@ -172,13 +177,13 @@ export const refineOptions: Record<
       function_call: {
         name: "extract_sentiment",
       },
-    };
+    });
 
     Here's another example of adding a function:
 
     Before:
 
-    prompt = {
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-3.5-turbo",
         messages: [
           {
@@ -187,29 +192,29 @@ export const refineOptions: Record<
 
             title: \${scenario.title}
             body: \${scenario.body}
-            
+
             On a scale from 1 to 3, how likely is it that the person writing this post has the following need? If you are not sure, make your best guess, or answer 1.
-            
+
             Need: \${scenario.need}
-            
+
             Answer one integer between 1 and 3.\`,
           },
         ],
         temperature: 0,
-    };
+    });
 
     After:
 
-    prompt = {
+    definePrompt("openai/ChatCompletion", {
         model: "gpt-3.5-turbo",
         messages: [
           {
             role: "user",
             content: \`Title: \${scenario.title}
       Body: \${scenario.body}
-      
+
       Need: \${scenario.need}
-      
+
       Rate likelihood on 1-3 scale.\`,
           },
         ],
@@ -230,8 +235,53 @@ export const refineOptions: Record<
         function_call: {
           name: "score_post",
         },
-      };
-    
+      });
+
+    Another example
+
+    Before:
+
+    definePrompt("openai/ChatCompletion", {
+      model: "gpt-3.5-turbo",
+      stream: true,
+      messages: [
+        {
+          role: "system",
+          content: \`Write 'Start experimenting!' in \${scenario.language}\`,
+        },
+      ],
+    });
+
+    After:
+
+    definePrompt("openai/ChatCompletion", {
+      model: "gpt-3.5-turbo",
+      messages: [
+        {
+          role: "system",
+          content: \`Write 'Start experimenting!' in \${scenario.language}\`,
+        },
+      ],
+      functions: [
+        {
+          name: "write_in_language",
+          parameters: {
+            type: "object",
+            properties: {
+              text: {
+                type: "string",
+              },
+            },
+          },
+        },
+      ],
+      function_call: {
+        name: "write_in_language",
+      },
+    });
+
     Add an OpenAI function that takes one or more nested parameters that match the expected output from this prompt.`,
+    },
   },
+  "replicate/llama2": {},
 };

src/components/SelectModelModal/SelectModelModal.tsx

@@ -1,85 +0,0 @@
-import {
-  Button,
-  Modal,
-  ModalBody,
-  ModalCloseButton,
-  ModalContent,
-  ModalFooter,
-  ModalHeader,
-  ModalOverlay,
-  VStack,
-  Text,
-  Spinner,
-  HStack,
-  Icon,
-} from "@chakra-ui/react";
-import { RiExchangeFundsFill } from "react-icons/ri";
-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 createMutation = api.promptVariants.create.useMutation();
-
-  const [createNewVariant, creationInProgress] = useHandledAsyncCallback(async () => {
-    if (!experiment?.data?.id) return;
-    await createMutation.mutateAsync({
-      experimentId: experiment?.data?.id,
-      variantId,
-      newModel: selectedModel,
-    });
-    await utils.promptVariants.list.invalidate();
-    onClose();
-  }, [createMutation, experiment?.data?.id, variantId, onClose]);
-
-  return (
-    <Modal isOpen onClose={onClose} size={{ base: "xl", sm: "2xl", md: "3xl" }}>
-      <ModalOverlay />
-      <ModalContent w={1200}>
-        <ModalHeader>
-          <HStack>
-            <Icon as={RiExchangeFundsFill} />
-            <Text>Change Model</Text>
-          </HStack>
-        </ModalHeader>
-        <ModalCloseButton />
-        <ModalBody maxW="unset">
-          <VStack spacing={8}>
-            <ModelStatsCard label="Original Model" model={originalModel} />
-            {originalModel !== selectedModel && (
-              <ModelStatsCard label="New Model" model={selectedModel} />
-            )}
-            <SelectModelSearch selectedModel={selectedModel} setSelectedModel={setSelectedModel} />
-          </VStack>
-        </ModalBody>
-
-        <ModalFooter>
-          <Button
-            colorScheme="blue"
-            onClick={createNewVariant}
-            minW={24}
-            disabled={originalModel === selectedModel}
-          >
-            {creationInProgress ? <Spinner boxSize={4} /> : <Text>Continue</Text>}
-          </Button>
-        </ModalFooter>
-      </ModalContent>
-    </Modal>
-  );
-};

src/components/SelectModelModal/SelectModelSearch.tsx

@@ -1,47 +0,0 @@
-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>
-  );
-};

src/components/VariantHeader/VariantHeaderMenuButton.tsx

@@ -17,8 +17,7 @@ import { FaRegClone } from "react-icons/fa";
 import { useState } from "react";
 import { RefinePromptModal } from "../RefinePromptModal/RefinePromptModal";
 import { RiExchangeFundsFill } from "react-icons/ri";
-import { SelectModelModal } from "../SelectModelModal/SelectModelModal";
-import { type SupportedModel } from "~/server/types";
+import { ChangeModelModal } from "../ChangeModelModal/ChangeModelModal";
 
 export default function VariantHeaderMenuButton({
   variant,
@@ -51,7 +50,7 @@ export default function VariantHeaderMenuButton({
     await utils.promptVariants.list.invalidate();
   }, [hideMutation, variant.id]);
 
-  const [selectModelModalOpen, setSelectModelModalOpen] = useState(false);
+  const [changeModelModalOpen, setChangeModelModalOpen] = useState(false);
   const [refinePromptModalOpen, setRefinePromptModalOpen] = useState(false);
 
   return (
@@ -73,7 +72,7 @@ export default function VariantHeaderMenuButton({
           </MenuItem>
           <MenuItem
             icon={<Icon as={RiExchangeFundsFill} boxSize={5} />}
-            onClick={() => setSelectModelModalOpen(true)}
+            onClick={() => setChangeModelModalOpen(true)}
           >
             Change Model
           </MenuItem>
@@ -98,12 +97,8 @@ export default function VariantHeaderMenuButton({
           )}
         </MenuList>
       </Menu>
-      {selectModelModalOpen && (
-        <SelectModelModal
-          originalModel={variant.model as SupportedModel}
-          variantId={variant.id}
-          onClose={() => setSelectModelModalOpen(false)}
-        />
+      {changeModelModalOpen && (
+        <ChangeModelModal variant={variant} onClose={() => setChangeModelModalOpen(false)} />
       )}
       {refinePromptModalOpen && (
         <RefinePromptModal variant={variant} onClose={() => setRefinePromptModalOpen(false)} />

src/env.mjs

@@ -9,7 +9,6 @@ export const env = createEnv({
   server: {
     DATABASE_URL: z.string().url(),
     NODE_ENV: z.enum(["development", "test", "production"]).default("development"),
-    OPENAI_API_KEY: z.string().min(1),
     RESTRICT_PRISMA_LOGS: z
       .string()
       .optional()
@@ -17,7 +16,8 @@ export const env = createEnv({
       .transform((val) => val.toLowerCase() === "true"),
     GITHUB_CLIENT_ID: z.string().min(1),
     GITHUB_CLIENT_SECRET: z.string().min(1),
-    REPLICATE_API_TOKEN: z.string().min(1),
+    OPENAI_API_KEY: z.string().min(1),
+    REPLICATE_API_TOKEN: z.string().default("placeholder"),
   },
 
   /**

src/modelProviders/frontendModelProviders.ts

@@ -1,14 +1,15 @@
 import openaiChatCompletionFrontend from "./openai-ChatCompletion/frontend";
 import replicateLlama2Frontend from "./replicate-llama2/frontend";
+import { type SupportedProvider, type FrontendModelProvider } from "./types";
 
 // TODO: make sure we get a typescript error if you forget to add a provider here
 
 // Keep attributes here that need to be accessible from the frontend. We can't
 // just include them in the default `modelProviders` object because it has some
 // transient dependencies that can only be imported on the server.
-const modelProvidersFrontend = {
+const frontendModelProviders: Record<SupportedProvider, FrontendModelProvider<any, any>> = {
   "openai/ChatCompletion": openaiChatCompletionFrontend,
   "replicate/llama2": replicateLlama2Frontend,
-} as const;
+};
 
-export default modelProvidersFrontend;
+export default frontendModelProviders;

src/modelProviders/modelProviders.ts

@@ -1,9 +1,10 @@
 import openaiChatCompletion from "./openai-ChatCompletion";
 import replicateLlama2 from "./replicate-llama2";
+import { type SupportedProvider, type ModelProvider } from "./types";
 
-const modelProviders = {
+const modelProviders: Record<SupportedProvider, ModelProvider<any, any, any>> = {
   "openai/ChatCompletion": openaiChatCompletion,
   "replicate/llama2": replicateLlama2,
-} as const;
+};
 
 export default modelProviders;

src/modelProviders/modelStats.ts

@@ -1,77 +0,0 @@
-import { type SupportedModel } from "../server/types";
-
-interface ModelStats {
-  contextLength: number;
-  promptTokenPrice: number;
-  completionTokenPrice: number;
-  speed: "fast" | "medium" | "slow";
-  provider: "OpenAI";
-  learnMoreUrl: string;
-}
-
-export const modelStats: Record<SupportedModel, ModelStats> = {
-  "gpt-4": {
-    contextLength: 8192,
-    promptTokenPrice: 0.00003,
-    completionTokenPrice: 0.00006,
-    speed: "medium",
-    provider: "OpenAI",
-    learnMoreUrl: "https://openai.com/gpt-4",
-  },
-  "gpt-4-0613": {
-    contextLength: 8192,
-    promptTokenPrice: 0.00003,
-    completionTokenPrice: 0.00006,
-    speed: "medium",
-    provider: "OpenAI",
-    learnMoreUrl: "https://openai.com/gpt-4",
-  },
-  "gpt-4-32k": {
-    contextLength: 32768,
-    promptTokenPrice: 0.00006,
-    completionTokenPrice: 0.00012,
-    speed: "medium",
-    provider: "OpenAI",
-    learnMoreUrl: "https://openai.com/gpt-4",
-  },
-  "gpt-4-32k-0613": {
-    contextLength: 32768,
-    promptTokenPrice: 0.00006,
-    completionTokenPrice: 0.00012,
-    speed: "medium",
-    provider: "OpenAI",
-    learnMoreUrl: "https://openai.com/gpt-4",
-  },
-  "gpt-3.5-turbo": {
-    contextLength: 4096,
-    promptTokenPrice: 0.0000015,
-    completionTokenPrice: 0.000002,
-    speed: "fast",
-    provider: "OpenAI",
-    learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-  },
-  "gpt-3.5-turbo-0613": {
-    contextLength: 4096,
-    promptTokenPrice: 0.0000015,
-    completionTokenPrice: 0.000002,
-    speed: "fast",
-    provider: "OpenAI",
-    learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-  },
-  "gpt-3.5-turbo-16k": {
-    contextLength: 16384,
-    promptTokenPrice: 0.000003,
-    completionTokenPrice: 0.000004,
-    speed: "fast",
-    provider: "OpenAI",
-    learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-  },
-  "gpt-3.5-turbo-16k-0613": {
-    contextLength: 16384,
-    promptTokenPrice: 0.000003,
-    completionTokenPrice: 0.000004,
-    speed: "fast",
-    provider: "OpenAI",
-    learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-  },
-};

src/modelProviders/openai-ChatCompletion/codegen/codegen.ts

@@ -56,6 +56,14 @@ modelProperty.type = "string";
 modelProperty.enum = modelProperty.oneOf[1].enum;
 delete modelProperty["oneOf"];
 
+// The default of "inf" confuses the Typescript generator, so can just remove it
+assert(
+  "max_tokens" in completionRequestSchema.properties &&
+    isObject(completionRequestSchema.properties.max_tokens) &&
+    "default" in completionRequestSchema.properties.max_tokens,
+);
+delete completionRequestSchema.properties.max_tokens["default"];
+
 // Get the directory of the current script
 const currentDirectory = path.dirname(import.meta.url).replace("file://", "");
 

src/modelProviders/openai-ChatCompletion/codegen/input.schema.json

@@ -150,7 +150,6 @@
     },
     "max_tokens": {
       "description": "The maximum number of [tokens](/tokenizer) to generate in the chat completion.\n\nThe total length of input tokens and generated tokens is limited by the model's context length. [Example Python code](https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb) for counting tokens.\n",
-      "default": "inf",
       "type": "integer"
     },
     "presence_penalty": {

src/modelProviders/openai-ChatCompletion/frontend.ts

@@ -1,8 +1,50 @@
 import { type JsonValue } from "type-fest";
-import { type OpenaiChatModelProvider } from ".";
-import { type ModelProviderFrontend } from "../types";
+import { type SupportedModel } from ".";
+import { type FrontendModelProvider } from "../types";
+import { type ChatCompletion } from "openai/resources/chat";
+
+const frontendModelProvider: FrontendModelProvider<SupportedModel, ChatCompletion> = {
+  name: "OpenAI ChatCompletion",
+
+  models: {
+    "gpt-4-0613": {
+      name: "GPT-4",
+      contextWindow: 8192,
+      promptTokenPrice: 0.00003,
+      completionTokenPrice: 0.00006,
+      speed: "medium",
+      provider: "openai/ChatCompletion",
+      learnMoreUrl: "https://openai.com/gpt-4",
+    },
+    "gpt-4-32k-0613": {
+      name: "GPT-4 32k",
+      contextWindow: 32768,
+      promptTokenPrice: 0.00006,
+      completionTokenPrice: 0.00012,
+      speed: "medium",
+      provider: "openai/ChatCompletion",
+      learnMoreUrl: "https://openai.com/gpt-4",
+    },
+    "gpt-3.5-turbo-0613": {
+      name: "GPT-3.5 Turbo",
+      contextWindow: 4096,
+      promptTokenPrice: 0.0000015,
+      completionTokenPrice: 0.000002,
+      speed: "fast",
+      provider: "openai/ChatCompletion",
+      learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
+    },
+    "gpt-3.5-turbo-16k-0613": {
+      name: "GPT-3.5 Turbo 16k",
+      contextWindow: 16384,
+      promptTokenPrice: 0.000003,
+      completionTokenPrice: 0.000004,
+      speed: "fast",
+      provider: "openai/ChatCompletion",
+      learnMoreUrl: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
+    },
+  },
 
-const modelProviderFrontend: ModelProviderFrontend<OpenaiChatModelProvider> = {
   normalizeOutput: (output) => {
     const message = output.choices[0]?.message;
     if (!message)
@@ -39,4 +81,4 @@ const modelProviderFrontend: ModelProviderFrontend<OpenaiChatModelProvider> = {
   },
 };
 
-export default modelProviderFrontend;
+export default frontendModelProvider;

src/modelProviders/openai-ChatCompletion/getCompletion.ts

@@ -8,10 +8,10 @@ import { countOpenAIChatTokens } from "~/utils/countTokens";
 import { type CompletionResponse } from "../types";
 import { omit } from "lodash-es";
 import { openai } from "~/server/utils/openai";
-import { type OpenAIChatModel } from "~/server/types";
 import { truthyFilter } from "~/utils/utils";
 import { APIError } from "openai";
-import { modelStats } from "../modelStats";
+import frontendModelProvider from "./frontend";
+import modelProvider, { type SupportedModel } from ".";
 
 const mergeStreamedChunks = (
   base: ChatCompletion | null,
@@ -60,6 +60,7 @@ export async function getCompletion(
   let finalCompletion: ChatCompletion | null = null;
   let promptTokens: number | undefined = undefined;
   let completionTokens: number | undefined = undefined;
+  const modelName = modelProvider.getModel(input) as SupportedModel;
 
   try {
     if (onStream) {
@@ -81,12 +82,9 @@ export async function getCompletion(
         };
       }
       try {
-        promptTokens = countOpenAIChatTokens(
-          input.model as keyof typeof OpenAIChatModel,
-          input.messages,
-        );
+        promptTokens = countOpenAIChatTokens(modelName, input.messages);
         completionTokens = countOpenAIChatTokens(
-          input.model as keyof typeof OpenAIChatModel,
+          modelName,
           finalCompletion.choices.map((c) => c.message).filter(truthyFilter),
         );
       } catch (err) {
@@ -106,10 +104,10 @@ export async function getCompletion(
     }
     const timeToComplete = Date.now() - start;
 
-    const stats = modelStats[input.model as keyof typeof OpenAIChatModel];
+    const { promptTokenPrice, completionTokenPrice } = frontendModelProvider.models[modelName];
     let cost = undefined;
-    if (stats && promptTokens && completionTokens) {
-      cost = promptTokens * stats.promptTokenPrice + completionTokens * stats.completionTokenPrice;
+    if (promptTokenPrice && completionTokenPrice && promptTokens && completionTokens) {
+      cost = promptTokens * promptTokenPrice + completionTokens * completionTokenPrice;
     }
 
     return {

src/modelProviders/openai-ChatCompletion/index.ts

@@ -3,6 +3,7 @@ import { type ModelProvider } from "../types";
 import inputSchema from "./codegen/input.schema.json";
 import { type ChatCompletion, type CompletionCreateParams } from "openai/resources/chat";
 import { getCompletion } from "./getCompletion";
+import frontendModelProvider from "./frontend";
 
 const supportedModels = [
   "gpt-4-0613",
@@ -11,7 +12,7 @@ const supportedModels = [
   "gpt-3.5-turbo-16k-0613",
 ] as const;
 
-type SupportedModel = (typeof supportedModels)[number];
+export type SupportedModel = (typeof supportedModels)[number];
 
 export type OpenaiChatModelProvider = ModelProvider<
   SupportedModel,
@@ -20,25 +21,6 @@ export type OpenaiChatModelProvider = ModelProvider<
 >;
 
 const modelProvider: OpenaiChatModelProvider = {
-  name: "OpenAI ChatCompletion",
-  models: {
-    "gpt-4-0613": {
-      name: "GPT-4",
-      learnMore: "https://openai.com/gpt-4",
-    },
-    "gpt-4-32k-0613": {
-      name: "GPT-4 32k",
-      learnMore: "https://openai.com/gpt-4",
-    },
-    "gpt-3.5-turbo-0613": {
-      name: "GPT-3.5 Turbo",
-      learnMore: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-    },
-    "gpt-3.5-turbo-16k-0613": {
-      name: "GPT-3.5 Turbo 16k",
-      learnMore: "https://platform.openai.com/docs/guides/gpt/chat-completions-api",
-    },
-  },
   getModel: (input) => {
     if (supportedModels.includes(input.model as SupportedModel))
       return input.model as SupportedModel;
@@ -57,6 +39,7 @@ const modelProvider: OpenaiChatModelProvider = {
   inputSchema: inputSchema as JSONSchema4,
   shouldStream: (input) => input.stream ?? false,
   getCompletion,
+  ...frontendModelProvider,
 };
 
 export default modelProvider;

src/modelProviders/replicate-llama2/frontend.ts

@@ -1,7 +1,36 @@
-import { type ReplicateLlama2Provider } from ".";
-import { type ModelProviderFrontend } from "../types";
+import { type SupportedModel, type ReplicateLlama2Output } from ".";
+import { type FrontendModelProvider } from "../types";
+
+const frontendModelProvider: FrontendModelProvider<SupportedModel, ReplicateLlama2Output> = {
+  name: "Replicate Llama2",
+
+  models: {
+    "7b-chat": {
+      name: "LLama 2 7B Chat",
+      contextWindow: 4096,
+      pricePerSecond: 0.0023,
+      speed: "fast",
+      provider: "replicate/llama2",
+      learnMoreUrl: "https://replicate.com/a16z-infra/llama7b-v2-chat",
+    },
+    "13b-chat": {
+      name: "LLama 2 13B Chat",
+      contextWindow: 4096,
+      pricePerSecond: 0.0023,
+      speed: "medium",
+      provider: "replicate/llama2",
+      learnMoreUrl: "https://replicate.com/a16z-infra/llama13b-v2-chat",
+    },
+    "70b-chat": {
+      name: "LLama 2 70B Chat",
+      contextWindow: 4096,
+      pricePerSecond: 0.0032,
+      speed: "slow",
+      provider: "replicate/llama2",
+      learnMoreUrl: "https://replicate.com/replicate/llama70b-v2-chat",
+    },
+  },
 
-const modelProviderFrontend: ModelProviderFrontend<ReplicateLlama2Provider> = {
   normalizeOutput: (output) => {
     return {
       type: "text",
@@ -10,4 +39,4 @@ const modelProviderFrontend: ModelProviderFrontend<ReplicateLlama2Provider> = {
   },
 };
 
-export default modelProviderFrontend;
+export default frontendModelProvider;

src/modelProviders/replicate-llama2/index.ts

@@ -1,9 +1,10 @@
 import { type ModelProvider } from "../types";
+import frontendModelProvider from "./frontend";
 import { getCompletion } from "./getCompletion";
 
 const supportedModels = ["7b-chat", "13b-chat", "70b-chat"] as const;
 
-type SupportedModel = (typeof supportedModels)[number];
+export type SupportedModel = (typeof supportedModels)[number];
 
 export type ReplicateLlama2Input = {
   model: SupportedModel;
@@ -25,12 +26,6 @@ export type ReplicateLlama2Provider = ModelProvider<
 >;
 
 const modelProvider: ReplicateLlama2Provider = {
-  name: "OpenAI ChatCompletion",
-  models: {
-    "7b-chat": {},
-    "13b-chat": {},
-    "70b-chat": {},
-  },
   getModel: (input) => {
     if (supportedModels.includes(input.model)) return input.model;
 
@@ -69,6 +64,7 @@ const modelProvider: ReplicateLlama2Provider = {
   },
   shouldStream: (input) => input.stream ?? false,
   getCompletion,
+  ...frontendModelProvider,
 };
 
 export default modelProvider;

src/modelProviders/types.ts

@@ -1,9 +1,33 @@
 import { type JSONSchema4 } from "json-schema";
 import { type JsonValue } from "type-fest";
+import { z } from "zod";
 
-type ModelProviderModel = {
-  name?: string;
-  learnMore?: string;
+const ZodSupportedProvider = z.union([
+  z.literal("openai/ChatCompletion"),
+  z.literal("replicate/llama2"),
+]);
+
+export type SupportedProvider = z.infer<typeof ZodSupportedProvider>;
+
+export const ZodModel = z.object({
+  name: z.string(),
+  contextWindow: z.number(),
+  promptTokenPrice: z.number().optional(),
+  completionTokenPrice: z.number().optional(),
+  pricePerSecond: z.number().optional(),
+  speed: z.union([z.literal("fast"), z.literal("medium"), z.literal("slow")]),
+  provider: ZodSupportedProvider,
+  description: z.string().optional(),
+  learnMoreUrl: z.string().optional(),
+});
+
+export type Model = z.infer<typeof ZodModel>;
+
+export type FrontendModelProvider<SupportedModels extends string, OutputSchema> = {
+  name: string;
+  models: Record<SupportedModels, Model>;
+
+  normalizeOutput: (output: OutputSchema) => NormalizedOutput;
 };
 
 export type CompletionResponse<T> =
@@ -19,8 +43,6 @@ export type CompletionResponse<T> =
     };
 
 export type ModelProvider<SupportedModels extends string, InputSchema, OutputSchema> = {
-  name: string;
-  models: Record<SupportedModels, ModelProviderModel>;
   getModel: (input: InputSchema) => SupportedModels | null;
   shouldStream: (input: InputSchema) => boolean;
   inputSchema: JSONSchema4;
@@ -31,7 +53,7 @@ export type ModelProvider<SupportedModels extends string, InputSchema, OutputSch
 
   // This is just a convenience for type inference, don't use it at runtime
   _outputSchema?: OutputSchema | null;
-};
+} & FrontendModelProvider<SupportedModels, OutputSchema>;
 
 export type NormalizedOutput =
   | {
@@ -42,7 +64,3 @@ export type NormalizedOutput =
       type: "json";
       value: JsonValue;
     };
-
-export type ModelProviderFrontend<ModelProviderT extends ModelProvider<any, any, any>> = {
-  normalizeOutput: (output: NonNullable<ModelProviderT["_outputSchema"]>) => NormalizedOutput;
-};

src/pages/experiments/index.tsx

@@ -20,22 +20,25 @@ export default function ExperimentsPage() {
   const experiments = api.experiments.list.useQuery();
 
   const user = useSession().data;
+  const authLoading = useSession().status === "loading";
 
-  if (user === null) {
+  if (user === null || authLoading) {
     return (
       <AppShell title="Experiments">
         <Center h="100%">
-          <Text>
-            <Link
-              onClick={() => {
-                signIn("github").catch(console.error);
-              }}
-              textDecor="underline"
-            >
-              Sign in
-            </Link>{" "}
-            to view or create new experiments!
-          </Text>
+          {!authLoading && (
+            <Text>
+              <Link
+                onClick={() => {
+                  signIn("github").catch(console.error);
+                }}
+                textDecor="underline"
+              >
+                Sign in
+              </Link>{" "}
+              to view or create new experiments!
+            </Text>
+          )}
         </Center>
       </AppShell>
     );

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

@@ -98,7 +98,7 @@ export const experimentsRouter = createTRPCRouter({
       },
     });
 
-    const [variant, _, scenario] = await prisma.$transaction([
+    const [variant, _, scenario1, scenario2, scenario3] = await prisma.$transaction([
       prisma.promptVariant.create({
         data: {
           experimentId: exp.id,
@@ -121,7 +121,7 @@ export const experimentsRouter = createTRPCRouter({
             messages: [
               {
                 role: "system",
-                content: \`"Return 'this is output for the scenario "${"$"}{scenario.text}"'\`,
+                content: \`Write 'Start experimenting!' in ${"$"}{scenario.language}\`,
               },
             ],
           });`,
@@ -133,20 +133,38 @@ export const experimentsRouter = createTRPCRouter({
       prisma.templateVariable.create({
         data: {
           experimentId: exp.id,
-          label: "text",
+          label: "language",
         },
       }),
       prisma.testScenario.create({
         data: {
           experimentId: exp.id,
           variableValues: {
-            text: "This is a test scenario.",
+            language: "English",
+          },
+        },
+      }),
+      prisma.testScenario.create({
+        data: {
+          experimentId: exp.id,
+          variableValues: {
+            language: "Spanish",
+          },
+        },
+      }),
+      prisma.testScenario.create({
+        data: {
+          experimentId: exp.id,
+          variableValues: {
+            language: "German",
           },
         },
       }),
     ]);
 
-    await generateNewCell(variant.id, scenario.id);
+    await generateNewCell(variant.id, scenario1.id);
+    await generateNewCell(variant.id, scenario2.id);
+    await generateNewCell(variant.id, scenario3.id);
 
     return exp;
   }),

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

@@ -2,7 +2,6 @@ import { z } from "zod";
 import { createTRPCRouter, protectedProcedure, publicProcedure } from "~/server/api/trpc";
 import { prisma } from "~/server/db";
 import { generateNewCell } from "~/server/utils/generateNewCell";
-import { type SupportedModel } from "~/server/types";
 import userError from "~/server/utils/error";
 import { recordExperimentUpdated } from "~/server/utils/recordExperimentUpdated";
 import { reorderPromptVariants } from "~/server/utils/reorderPromptVariants";
@@ -10,6 +9,7 @@ import { type PromptVariant } from "@prisma/client";
 import { deriveNewConstructFn } from "~/server/utils/deriveNewContructFn";
 import { requireCanModifyExperiment, requireCanViewExperiment } from "~/utils/accessControl";
 import parseConstructFn from "~/server/utils/parseConstructFn";
+import { ZodModel } from "~/modelProviders/types";
 
 export const promptVariantsRouter = createTRPCRouter({
   list: publicProcedure
@@ -144,7 +144,7 @@ export const promptVariantsRouter = createTRPCRouter({
       z.object({
         experimentId: z.string(),
         variantId: z.string().optional(),
-        newModel: z.string().optional(),
+        newModel: ZodModel.optional(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
@@ -186,10 +186,7 @@ export const promptVariantsRouter = createTRPCRouter({
           ? `${originalVariant?.label} Copy`
           : `Prompt Variant ${largestSortIndex + 2}`;
 
-      const newConstructFn = await deriveNewConstructFn(
-        originalVariant,
-        input.newModel as SupportedModel,
-      );
+      const newConstructFn = await deriveNewConstructFn(originalVariant, input.newModel);
 
       const createNewVariantAction = prisma.promptVariant.create({
         data: {
@@ -284,11 +281,12 @@ export const promptVariantsRouter = createTRPCRouter({
       return updatedPromptVariant;
     }),
 
-  getRefinedPromptFn: protectedProcedure
+  getModifiedPromptFn: protectedProcedure
     .input(
       z.object({
         id: z.string(),
-        instructions: z.string(),
+        instructions: z.string().optional(),
+        newModel: ZodModel.optional(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
@@ -307,7 +305,7 @@ export const promptVariantsRouter = createTRPCRouter({
 
       const promptConstructionFn = await deriveNewConstructFn(
         existing,
-        constructedPrompt.model as SupportedModel,
+        input.newModel,
         input.instructions,
       );
 

src/server/tasks/queryLLM.task.ts

@@ -99,7 +99,6 @@ export const queryLLM = defineTask<queryLLMJob>("queryLLM", async (task) => {
 
   const provider = modelProviders[prompt.modelProvider];
 
-  // @ts-expect-error TODO FIX ASAP
   const streamingChannel = provider.shouldStream(prompt.modelInput) ? generateChannel() : null;
 
   if (streamingChannel) {
@@ -116,8 +115,6 @@ export const queryLLM = defineTask<queryLLMJob>("queryLLM", async (task) => {
     : null;
 
   for (let i = 0; true; i++) {
-    // @ts-expect-error TODO FIX ASAP
-
     const response = await provider.getCompletion(prompt.modelInput, onStream);
     if (response.type === "success") {
       const inputHash = hashPrompt(prompt);
@@ -126,7 +123,7 @@ export const queryLLM = defineTask<queryLLMJob>("queryLLM", async (task) => {
         data: {
           scenarioVariantCellId,
           inputHash,
-          output: response.value as unknown as Prisma.InputJsonObject,
+          output: response.value as Prisma.InputJsonObject,
           timeToComplete: response.timeToComplete,
           promptTokens: response.promptTokens,
           completionTokens: response.completionTokens,
@@ -154,7 +151,7 @@ export const queryLLM = defineTask<queryLLMJob>("queryLLM", async (task) => {
           errorMessage: response.message,
           statusCode: response.statusCode,
           retryTime: shouldRetry ? new Date(Date.now() + delay) : null,
-          retrievalStatus: shouldRetry ? "PENDING" : "ERROR",
+          retrievalStatus: "ERROR",
         },
       });
 

src/server/types.ts

@@ -1,12 +0,0 @@
-export enum OpenAIChatModel {
-  "gpt-4" = "gpt-4",
-  "gpt-4-0613" = "gpt-4-0613",
-  "gpt-4-32k" = "gpt-4-32k",
-  "gpt-4-32k-0613" = "gpt-4-32k-0613",
-  "gpt-3.5-turbo" = "gpt-3.5-turbo",
-  "gpt-3.5-turbo-0613" = "gpt-3.5-turbo-0613",
-  "gpt-3.5-turbo-16k" = "gpt-3.5-turbo-16k",
-  "gpt-3.5-turbo-16k-0613" = "gpt-3.5-turbo-16k-0613",
-}
-
-export type SupportedModel = keyof typeof OpenAIChatModel;

src/server/utils/deriveNewContructFn.ts

@@ -1,18 +1,18 @@
 import { type PromptVariant } from "@prisma/client";
-import { type SupportedModel } from "../types";
 import ivm from "isolated-vm";
 import dedent from "dedent";
 import { openai } from "./openai";
-import { getApiShapeForModel } from "./getTypesForModel";
 import { isObject } from "lodash-es";
 import { type CompletionCreateParams } from "openai/resources/chat/completions";
 import formatPromptConstructor from "~/utils/formatPromptConstructor";
+import { type SupportedProvider, type Model } from "~/modelProviders/types";
+import modelProviders from "~/modelProviders/modelProviders";
 
 const isolate = new ivm.Isolate({ memoryLimit: 128 });
 
 export async function deriveNewConstructFn(
   originalVariant: PromptVariant | null,
-  newModel?: SupportedModel,
+  newModel?: Model,
   instructions?: string,
 ) {
   if (originalVariant && !newModel && !instructions) {
@@ -36,10 +36,11 @@ export async function deriveNewConstructFn(
 const NUM_RETRIES = 5;
 const requestUpdatedPromptFunction = async (
   originalVariant: PromptVariant,
-  newModel?: SupportedModel,
+  newModel?: Model,
   instructions?: string,
 ) => {
-  const originalModel = originalVariant.model as SupportedModel;
+  const originalModelProvider = modelProviders[originalVariant.modelProvider as SupportedProvider];
+  const originalModel = originalModelProvider.models[originalVariant.model] as Model;
   let newContructionFn = "";
   for (let i = 0; i < NUM_RETRIES; i++) {
     try {
@@ -47,17 +48,33 @@ const requestUpdatedPromptFunction = async (
         {
           role: "system",
           content: `Your job is to update prompt constructor functions. Here is the api shape for the current model:\n---\n${JSON.stringify(
-            getApiShapeForModel(originalModel),
+            originalModelProvider.inputSchema,
             null,
             2,
           )}\n\nDo not add any assistant messages.`,
         },
+        {
+          role: "user",
+          content: `This is the current prompt constructor function:\n---\n${originalVariant.constructFn}`,
+        },
       ];
       if (newModel) {
         messages.push({
           role: "user",
-          content: `Return the prompt constructor function for ${newModel} given the following prompt constructor function for ${originalModel}:\n---\n${originalVariant.constructFn}`,
+          content: `Return the prompt constructor function for ${newModel.name} given the existing prompt constructor function for ${originalModel.name}`,
         });
+        if (newModel.provider !== originalModel.provider) {
+          messages.push({
+            role: "user",
+            content: `The old provider was ${originalModel.provider}. The new provider is ${
+              newModel.provider
+            }. Here is the schema for the new model:\n---\n${JSON.stringify(
+              modelProviders[newModel.provider].inputSchema,
+              null,
+              2,
+            )}`,
+          });
+        }
       }
       if (instructions) {
         messages.push({
@@ -65,10 +82,6 @@ const requestUpdatedPromptFunction = async (
           content: instructions,
         });
       }
-      messages.push({
-        role: "system",
-        content: "The prompt variable has already been declared, so do not declare it again.",
-      });
       const completion = await openai.chat.completions.create({
         model: "gpt-4",
         messages,

src/server/utils/generateNewCell.ts

@@ -4,8 +4,9 @@ import { queueLLMRetrievalTask } from "./queueLLMRetrievalTask";
 import parseConstructFn from "./parseConstructFn";
 import { type JsonObject } from "type-fest";
 import hashPrompt from "./hashPrompt";
+import { omit } from "lodash-es";
 
-export const generateNewCell = async (variantId: string, scenarioId: string) => {
+export const generateNewCell = async (variantId: string, scenarioId: string): Promise<void> => {
   const variant = await prisma.promptVariant.findUnique({
     where: {
       id: variantId,
@@ -18,7 +19,7 @@ export const generateNewCell = async (variantId: string, scenarioId: string) =>
     },
   });
 
-  if (!variant || !scenario) return null;
+  if (!variant || !scenario) return;
 
   let cell = await prisma.scenarioVariantCell.findUnique({
     where: {
@@ -32,7 +33,7 @@ export const generateNewCell = async (variantId: string, scenarioId: string) =>
     },
   });
 
-  if (cell) return cell;
+  if (cell) return;
 
   const parsedConstructFn = await parseConstructFn(
     variant.constructFn,
@@ -40,7 +41,7 @@ export const generateNewCell = async (variantId: string, scenarioId: string) =>
   );
 
   if ("error" in parsedConstructFn) {
-    return await prisma.scenarioVariantCell.create({
+    await prisma.scenarioVariantCell.create({
       data: {
         promptVariantId: variantId,
         testScenarioId: scenarioId,
@@ -49,6 +50,7 @@ export const generateNewCell = async (variantId: string, scenarioId: string) =>
         retrievalStatus: "ERROR",
       },
     });
+    return;
   }
 
   const inputHash = hashPrompt(parsedConstructFn);
@@ -69,29 +71,33 @@ export const generateNewCell = async (variantId: string, scenarioId: string) =>
     where: { inputHash },
   });
 
-  let newModelOutput;
-
   if (matchingModelOutput) {
-    newModelOutput = await prisma.modelOutput.create({
+    const newModelOutput = await prisma.modelOutput.create({
       data: {
+        ...omit(matchingModelOutput, ["id"]),
         scenarioVariantCellId: cell.id,
-        inputHash,
         output: matchingModelOutput.output as Prisma.InputJsonValue,
-        timeToComplete: matchingModelOutput.timeToComplete,
-        cost: matchingModelOutput.cost,
-        promptTokens: matchingModelOutput.promptTokens,
-        completionTokens: matchingModelOutput.completionTokens,
-        createdAt: matchingModelOutput.createdAt,
-        updatedAt: matchingModelOutput.updatedAt,
       },
     });
     await prisma.scenarioVariantCell.update({
       where: { id: cell.id },
       data: { retrievalStatus: "COMPLETE" },
     });
+
+    // Copy over all eval results as well
+    await Promise.all(
+      (
+        await prisma.outputEvaluation.findMany({ where: { modelOutputId: matchingModelOutput.id } })
+      ).map(async (evaluation) => {
+        await prisma.outputEvaluation.create({
+          data: {
+            ...omit(evaluation, ["id"]),
+            modelOutputId: newModelOutput.id,
+          },
+        });
+      }),
+    );
   } else {
     cell = await queueLLMRetrievalTask(cell.id);
   }
-
-  return { ...cell, modelOutput: newModelOutput };
 };

src/server/utils/getTypesForModel.ts

@@ -1,6 +0,0 @@
-import { type SupportedModel } from "../types";
-
-export const getApiShapeForModel = (model: SupportedModel) => {
-  // if (model in OpenAIChatModel) return openAIChatApiShape;
-  return "";
-};

src/server/utils/parseConstructFn.ts

@@ -70,7 +70,6 @@ export default async function parseConstructFn(
     // We've validated the JSON schema so this should be safe
     const input = prompt.input as Parameters<(typeof provider)["getModel"]>[0];
 
-    // @ts-expect-error TODO FIX ASAP
     const model = provider.getModel(input);
     if (!model) {
       return {
@@ -80,8 +79,6 @@ export default async function parseConstructFn(
 
     return {
       modelProvider: prompt.modelProvider as keyof typeof modelProviders,
-      // @ts-expect-error TODO FIX ASAP
-
       model,
       modelInput: input,
     };

src/utils/countTokens.ts

@@ -1,6 +1,6 @@
 import { type ChatCompletion } from "openai/resources/chat";
 import { GPTTokens } from "gpt-tokens";
-import { type OpenAIChatModel } from "~/server/types";
+import { type SupportedModel } from "~/modelProviders/openai-ChatCompletion";
 
 interface GPTTokensMessageItem {
   name?: string;
@@ -9,7 +9,7 @@ interface GPTTokensMessageItem {
 }
 
 export const countOpenAIChatTokens = (
-  model: keyof typeof OpenAIChatModel,
+  model: SupportedModel,
   messages: ChatCompletion.Choice.Message[],
 ) => {
   return new GPTTokens({ model, messages: messages as unknown as GPTTokensMessageItem[] })

src/utils/utils.ts

@@ -1 +1,5 @@
+import { type Model } from "~/modelProviders/types";
+
 export const truthyFilter = <T>(x: T | null | undefined): x is T => Boolean(x);
+
+export const keyForModel = (model: Model) => `${model.provider}/${model.name}`;