Add request cost to OutputStats (#12)

src/components/OutputsTable/OutputCell.tsx

@@ -6,13 +6,16 @@ import SyntaxHighlighter from "react-syntax-highlighter";
 import { docco } from "react-syntax-highlighter/dist/cjs/styles/hljs";
 import stringify from "json-stringify-pretty-compact";
 import { useMemo, type ReactElement } from "react";
-import { BsCheck, BsClock, BsX } from "react-icons/bs";
+import { BsCheck, BsClock, BsX, BsCurrencyDollar } from "react-icons/bs";
 import { type ModelOutput } from "@prisma/client";
 import { type ChatCompletion } from "openai/resources/chat";
 import { generateChannel } from "~/utils/generateChannel";
 import { isObject } from "lodash";
 import useSocket from "~/utils/useSocket";
 import { evaluateOutput } from "~/server/utils/evaluateOutput";
+import { calculateTokenCost } from "~/utils/calculateTokenCost";
+import { type JSONSerializable, type SupportedModel } from "~/server/types";
+import { getModelName } from "~/server/utils/getModelName";
 
 export default function OutputCell({
   scenario,
@@ -37,6 +40,8 @@ export default function OutputCell({
   if (variant.config === null || Object.keys(variant.config).length === 0)
     disabledReason = "Save your prompt variant to see output";
 
+  const model = getModelName(variant.config as JSONSerializable);
+
   const shouldStream =
     isObject(variant) &&
     "config" in variant &&
@@ -110,7 +115,7 @@ export default function OutputCell({
             { maxLength: 40 }
           )}
         </SyntaxHighlighter>
-        <OutputStats modelOutput={output.data} scenario={scenario} />
+        <OutputStats model={model} modelOutput={output.data} scenario={scenario} />
       </Box>
     );
   }
@@ -121,15 +126,17 @@ export default function OutputCell({
   return (
     <Flex w="100%" h="100%" direction="column" justifyContent="space-between" whiteSpace="pre-wrap">
       {contentToDisplay}
-      {output.data && <OutputStats modelOutput={output.data} scenario={scenario} />}
+      {output.data && <OutputStats model={model} modelOutput={output.data} scenario={scenario} />}
     </Flex>
   );
 }
 
 const OutputStats = ({
+  model,
   modelOutput,
   scenario,
 }: {
+  model: SupportedModel | null;
   modelOutput: ModelOutput;
   scenario: Scenario;
 }) => {
@@ -138,6 +145,15 @@ const OutputStats = ({
   const evals =
     api.evaluations.list.useQuery({ experimentId: experiment.data?.id ?? "" }).data ?? [];
 
+  const promptTokens = modelOutput.promptTokens;
+  const completionTokens = modelOutput.completionTokens;
+
+  const promptCost = promptTokens && model ? calculateTokenCost(model, promptTokens) : 0;
+  const completionCost =
+    completionTokens && model ? calculateTokenCost(model, completionTokens, true) : 0;
+
+  const cost = promptCost + completionCost;
+
   return (
     <HStack align="center" color="gray.500" fontSize="xs" mt={2}>
       <HStack flex={1}>
@@ -155,8 +171,12 @@ const OutputStats = ({
           );
         })}
       </HStack>
-      <HStack>
-        <Icon as={BsClock} mr={0.5} />
+      <HStack spacing={0}>
+        <Icon as={BsCurrencyDollar} />
+        <Text mr={1}>{cost.toFixed(3)}</Text>
+      </HStack>
+      <HStack spacing={0.5}>
+        <Icon as={BsClock} />
         <Text>{(timeToComplete / 1000).toFixed(2)}s</Text>
       </HStack>
     </HStack>