From 9952dd93d87c53c4c972cbbc1b06bd1f96cc5a2a Mon Sep 17 00:00:00 2001
From: Kyle Corbitt <kyle@corbt.com>
Date: Mon, 24 Jul 2023 17:20:38 -0700
Subject: [PATCH] Only pass in model and provider

I got somewhat confused by the extra fields, sorry.

Also makes some frontend changes to track that state more directly although in retrospect not sure the frontend changes make things any better.
---
 .../ChangeModelModal/ChangeModelModal.tsx     | 59 ++++++++++---------
 .../ChangeModelModal/ModelSearch.tsx          | 52 ++++++----------
 .../ChangeModelModal/ModelStatsCard.tsx       | 21 ++++---
 src/modelProviders/types.ts                   | 26 ++++----
 .../api/routers/promptVariants.router.ts      | 23 +++++---
 src/utils/utils.ts                            | 11 +++-
 6 files changed, 101 insertions(+), 91 deletions(-)

diff --git a/src/components/ChangeModelModal/ChangeModelModal.tsx b/src/components/ChangeModelModal/ChangeModelModal.tsx
index 47e57c4..15ff507 100644
--- a/src/components/ChangeModelModal/ChangeModelModal.tsx
+++ b/src/components/ChangeModelModal/ChangeModelModal.tsx
@@ -1,5 +1,7 @@
 import {
   Button,
+  HStack,
+  Icon,
   Modal,
   ModalBody,
   ModalCloseButton,
@@ -7,24 +9,21 @@ import {
   ModalFooter,
   ModalHeader,
   ModalOverlay,
-  VStack,
-  Text,
   Spinner,
-  HStack,
-  Icon,
+  Text,
+  VStack,
 } 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";
+import { useState } from "react";
+import { RiExchangeFundsFill } from "react-icons/ri";
+import { type ProviderModel } from "~/modelProviders/types";
+import { api } from "~/utils/api";
+import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
+import { lookupModel, modelLabel } from "~/utils/utils";
+import CompareFunctions from "../RefinePromptModal/CompareFunctions";
+import { ModelSearch } from "./ModelSearch";
+import { ModelStatsCard } from "./ModelStatsCard";
 
 export const ChangeModelModal = ({
   variant,
@@ -33,11 +32,13 @@ export const ChangeModelModal = ({
   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 originalModel = lookupModel(variant.modelProvider, variant.model);
+  const [selectedModel, setSelectedModel] = useState({
+    provider: variant.modelProvider,
+    model: variant.model,
+  } as ProviderModel);
+  const [convertedModel, setConvertedModel] = useState<ProviderModel | undefined>();
+
   const utils = api.useContext();
 
   const experiment = useExperiment();
@@ -72,9 +73,10 @@ export const ChangeModelModal = ({
     onClose();
   }, [replaceVariantMutation, variant, onClose, modifiedPromptFn]);
 
-  const originalModelLabel = keyForModel(originalModel);
-  const selectedModelLabel = keyForModel(selectedModel);
-  const convertedModelLabel = convertedModel ? keyForModel(convertedModel) : undefined;
+  const originalLabel = modelLabel(variant.modelProvider, variant.model);
+  const selectedLabel = modelLabel(selectedModel.provider, selectedModel.model);
+  const convertedLabel =
+    convertedModel && modelLabel(convertedModel.provider, convertedModel.model);
 
   return (
     <Modal
@@ -94,16 +96,19 @@ export const ChangeModelModal = ({
         <ModalBody maxW="unset">
           <VStack spacing={8}>
             <ModelStatsCard label="Original Model" model={originalModel} />
-            {originalModelLabel !== selectedModelLabel && (
-              <ModelStatsCard label="New Model" model={selectedModel} />
+            {originalLabel !== selectedLabel && (
+              <ModelStatsCard
+                label="New Model"
+                model={lookupModel(selectedModel.provider, selectedModel.model)}
+              />
             )}
             <ModelSearch selectedModel={selectedModel} setSelectedModel={setSelectedModel} />
             {isString(modifiedPromptFn) && (
               <CompareFunctions
                 originalFunction={variant.constructFn}
                 newFunction={modifiedPromptFn}
-                leftTitle={originalModelLabel}
-                rightTitle={convertedModelLabel}
+                leftTitle={originalLabel}
+                rightTitle={convertedLabel}
               />
             )}
           </VStack>
@@ -115,7 +120,7 @@ export const ChangeModelModal = ({
               colorScheme="gray"
               onClick={getModifiedPromptFn}
               minW={24}
-              isDisabled={originalModel === selectedModel || modificationInProgress}
+              isDisabled={originalLabel === selectedLabel || modificationInProgress}
             >
               {modificationInProgress ? <Spinner boxSize={4} /> : <Text>Convert</Text>}
             </Button>
diff --git a/src/components/ChangeModelModal/ModelSearch.tsx b/src/components/ChangeModelModal/ModelSearch.tsx
index addfd5f..8f03de4 100644
--- a/src/components/ChangeModelModal/ModelSearch.tsx
+++ b/src/components/ChangeModelModal/ModelSearch.tsx
@@ -1,49 +1,35 @@
-import { VStack, Text } from "@chakra-ui/react";
-import { type LegacyRef, useCallback } from "react";
-import Select, { type SingleValue } from "react-select";
+import { Text, VStack } from "@chakra-ui/react";
+import { type LegacyRef } from "react";
+import Select from "react-select";
 import { useElementDimensions } from "~/utils/hooks";
 
+import { flatMap } from "lodash-es";
 import frontendModelProviders from "~/modelProviders/frontendModelProviders";
-import { type Model } from "~/modelProviders/types";
-import { keyForModel } from "~/utils/utils";
+import { type ProviderModel } from "~/modelProviders/types";
+import { modelLabel } from "~/utils/utils";
 
-const modelOptions: { label: string; value: Model }[] = [];
+const modelOptions = flatMap(Object.entries(frontendModelProviders), ([providerId, provider]) =>
+  Object.entries(provider.models).map(([modelId]) => ({
+    provider: providerId,
+    model: modelId,
+  })),
+) as ProviderModel[];
 
-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;
+export const ModelSearch = (props: {
+  selectedModel: ProviderModel;
+  setSelectedModel: (model: ProviderModel) => 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
+      <Select<ProviderModel>
         styles={{ control: (provided) => ({ ...provided, width: containerDimensions?.width }) }}
-        value={selectedOption}
+        getOptionLabel={(data) => modelLabel(data.provider, data.model)}
+        getOptionValue={(data) => modelLabel(data.provider, data.model)}
         options={modelOptions}
-        onChange={handleSelection}
+        onChange={(option) => option && props.setSelectedModel(option)}
       />
     </VStack>
   );
diff --git a/src/components/ChangeModelModal/ModelStatsCard.tsx b/src/components/ChangeModelModal/ModelStatsCard.tsx
index 33457a2..1ce42a3 100644
--- a/src/components/ChangeModelModal/ModelStatsCard.tsx
+++ b/src/components/ChangeModelModal/ModelStatsCard.tsx
@@ -1,15 +1,22 @@
 import {
-  VStack,
-  Text,
-  HStack,
-  type StackProps,
   GridItem,
-  SimpleGrid,
+  HStack,
   Link,
+  SimpleGrid,
+  Text,
+  VStack,
+  type StackProps,
 } from "@chakra-ui/react";
-import { type Model } from "~/modelProviders/types";
+import { type lookupModel } from "~/utils/utils";
 
-export const ModelStatsCard = ({ label, model }: { label: string; model: Model }) => {
+export const ModelStatsCard = ({
+  label,
+  model,
+}: {
+  label: string;
+  model: ReturnType<typeof lookupModel>;
+}) => {
+  if (!model) return null;
   return (
     <VStack w="full" align="start">
       <Text fontWeight="bold" fontSize="sm" textTransform="uppercase">
diff --git a/src/modelProviders/types.ts b/src/modelProviders/types.ts
index 18c5545..87111f2 100644
--- a/src/modelProviders/types.ts
+++ b/src/modelProviders/types.ts
@@ -3,26 +3,26 @@ import { type IconType } from "react-icons";
 import { type JsonValue } from "type-fest";
 import { z } from "zod";
 
-const ZodSupportedProvider = z.union([
+export 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 = {
+  name: string;
+  contextWindow: number;
+  promptTokenPrice?: number;
+  completionTokenPrice?: number;
+  pricePerSecond?: number;
+  speed: "fast" | "medium" | "slow";
+  provider: SupportedProvider;
+  description?: string;
+  learnMoreUrl?: string;
+};
 
-export type Model = z.infer<typeof ZodModel>;
+export type ProviderModel = { provider: z.infer<typeof ZodSupportedProvider>; model: string };
 
 export type RefinementAction = { icon?: IconType; description: string; instructions: string };
 
diff --git a/src/server/api/routers/promptVariants.router.ts b/src/server/api/routers/promptVariants.router.ts
index d88a53d..b6a32ef 100644
--- a/src/server/api/routers/promptVariants.router.ts
+++ b/src/server/api/routers/promptVariants.router.ts
@@ -9,7 +9,8 @@ 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";
+import modelProviders from "~/modelProviders/modelProviders";
+import { ZodSupportedProvider } from "~/modelProviders/types";
 
 export const promptVariantsRouter = createTRPCRouter({
   list: publicProcedure
@@ -144,7 +145,6 @@ export const promptVariantsRouter = createTRPCRouter({
       z.object({
         experimentId: z.string(),
         variantId: z.string().optional(),
-        newModel: ZodModel.optional(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
@@ -186,7 +186,7 @@ export const promptVariantsRouter = createTRPCRouter({
           ? `${originalVariant?.label} Copy`
           : `Prompt Variant ${largestSortIndex + 2}`;
 
-      const newConstructFn = await deriveNewConstructFn(originalVariant, input.newModel);
+      const newConstructFn = await deriveNewConstructFn(originalVariant);
 
       const createNewVariantAction = prisma.promptVariant.create({
         data: {
@@ -286,7 +286,12 @@ export const promptVariantsRouter = createTRPCRouter({
       z.object({
         id: z.string(),
         instructions: z.string().optional(),
-        newModel: ZodModel.optional(),
+        newModel: z
+          .object({
+            provider: ZodSupportedProvider,
+            model: z.string(),
+          })
+          .optional(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
@@ -303,11 +308,11 @@ export const promptVariantsRouter = createTRPCRouter({
         return userError(constructedPrompt.error);
       }
 
-      const promptConstructionFn = await deriveNewConstructFn(
-        existing,
-        input.newModel,
-        input.instructions,
-      );
+      const model = input.newModel
+        ? modelProviders[input.newModel.provider].models[input.newModel.model]
+        : undefined;
+
+      const promptConstructionFn = await deriveNewConstructFn(existing, model, input.instructions);
 
       // TODO: Validate promptConstructionFn
       // TODO: Record in some sort of history
diff --git a/src/utils/utils.ts b/src/utils/utils.ts
index 2644dfb..a1432b9 100644
--- a/src/utils/utils.ts
+++ b/src/utils/utils.ts
@@ -1,5 +1,12 @@
-import { type Model } from "~/modelProviders/types";
+import frontendModelProviders from "~/modelProviders/frontendModelProviders";
+import { type ProviderModel } 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}`;
+export const lookupModel = (provider: string, model: string) => {
+  const modelObj = frontendModelProviders[provider as ProviderModel["provider"]]?.models[model];
+  return modelObj ? { ...modelObj, provider } : null;
+};
+
+export const modelLabel = (provider: string, model: string) =>
+  `${provider}/${lookupModel(provider, model)?.name ?? model}`;