From d7fb4a72365a02a3f1cf1e87322c0175d9911981 Mon Sep 17 00:00:00 2001
From: jarenm <50216013+jarenm@users.noreply.github.com>
Date: Fri, 7 Jul 2023 20:34:45 -0600
Subject: [PATCH] add cost tooltip (#22)

* add cost tooltip

* add padding and style
---
 .../OutputsTable/OutputCell/OutputStats.tsx   | 11 ++--
 src/components/OutputsTable/VariantStats.tsx  | 26 +++++++--
 src/components/tooltip/CostTooltip.tsx        | 54 +++++++++++++++++++
 .../api/routers/promptVariants.router.ts      |  2 +-
 4 files changed, 83 insertions(+), 10 deletions(-)
 create mode 100644 src/components/tooltip/CostTooltip.tsx

diff --git a/src/components/OutputsTable/OutputCell/OutputStats.tsx b/src/components/OutputsTable/OutputCell/OutputStats.tsx
index 59bc17e..eebc882 100644
--- a/src/components/OutputsTable/OutputCell/OutputStats.tsx
+++ b/src/components/OutputsTable/OutputCell/OutputStats.tsx
@@ -7,6 +7,7 @@ import { calculateTokenCost } from "~/utils/calculateTokenCost";
 import { evaluateOutput } from "~/server/utils/evaluateOutput";
 import { HStack, Icon, Text } from "@chakra-ui/react";
 import { BsCheck, BsClock, BsCurrencyDollar, BsX } from "react-icons/bs";
+import { CostTooltip } from "~/components/tooltip/CostTooltip";
 
 export const OutputStats = ({
   model,
@@ -48,10 +49,12 @@ export const OutputStats = ({
           );
         })}
       </HStack>
-      <HStack spacing={0}>
-        <Icon as={BsCurrencyDollar} />
-        <Text mr={1}>{cost.toFixed(3)}</Text>
-      </HStack>
+      <CostTooltip promptTokens={promptTokens} completionTokens={completionTokens} cost={cost}>
+        <HStack spacing={0}>
+          <Icon as={BsCurrencyDollar} />
+          <Text mr={1}>{cost.toFixed(3)}</Text>
+        </HStack>
+      </CostTooltip>
       <HStack spacing={0.5}>
         <Icon as={BsClock} />
         <Text>{(timeToComplete / 1000).toFixed(2)}s</Text>
diff --git a/src/components/OutputsTable/VariantStats.tsx b/src/components/OutputsTable/VariantStats.tsx
index b0096e5..ee7f2d8 100644
--- a/src/components/OutputsTable/VariantStats.tsx
+++ b/src/components/OutputsTable/VariantStats.tsx
@@ -4,13 +4,23 @@ import { cellPadding } from "../constants";
 import { api } from "~/utils/api";
 import chroma from "chroma-js";
 import { BsCurrencyDollar } from "react-icons/bs";
+import { CostTooltip } from "../tooltip/CostTooltip";
 
 export default function VariantStats(props: { variant: PromptVariant }) {
   const { data } = api.promptVariants.stats.useQuery(
     {
       variantId: props.variant.id,
     },
-    { initialData: { evalResults: [], overallCost: 0, scenarioCount: 0, outputCount: 0 } }
+    {
+      initialData: {
+        evalResults: [],
+        overallCost: 0,
+        promptTokens: 0,
+        completionTokens: 0,
+        scenarioCount: 0,
+        outputCount: 0,
+      },
+    }
   );
 
   const [passColor, neutralColor, failColor] = useToken("colors", [
@@ -46,10 +56,16 @@ export default function VariantStats(props: { variant: PromptVariant }) {
         })}
       </HStack>
       {data.overallCost && (
-        <HStack spacing={0} align="center" color="gray.500" my="2">
-          <Icon as={BsCurrencyDollar} />
-          <Text mr={1}>{data.overallCost.toFixed(3)}</Text>
-        </HStack>
+        <CostTooltip
+          promptTokens={data.promptTokens}
+          completionTokens={data.completionTokens}
+          cost={data.overallCost}
+        >
+          <HStack spacing={0} align="center" color="gray.500" my="2">
+            <Icon as={BsCurrencyDollar} />
+            <Text mr={1}>{data.overallCost.toFixed(3)}</Text>
+          </HStack>
+        </CostTooltip>
       )}
     </HStack>
   );
diff --git a/src/components/tooltip/CostTooltip.tsx b/src/components/tooltip/CostTooltip.tsx
new file mode 100644
index 0000000..d6b05c0
--- /dev/null
+++ b/src/components/tooltip/CostTooltip.tsx
@@ -0,0 +1,54 @@
+import { HStack, Icon, Text, Tooltip, type TooltipProps, VStack, Divider } from "@chakra-ui/react";
+import { BsCurrencyDollar } from "react-icons/bs";
+
+type CostTooltipProps = {
+  promptTokens: number | null;
+  completionTokens: number | null;
+  cost: number;
+} & TooltipProps;
+
+export const CostTooltip = ({
+  promptTokens,
+  completionTokens,
+  cost,
+  children,
+  ...props
+}: CostTooltipProps) => {
+  return (
+    <Tooltip
+      borderRadius="8"
+      color="gray.800"
+      bgColor="gray.50"
+      py={2}
+      hasArrow
+      label={
+        <VStack fontSize="sm" w="200">
+          <VStack>
+            <Text fontWeight="bold">Cost</Text>
+            <HStack spacing={0}>
+              <Icon as={BsCurrencyDollar} />
+              <Text>{cost.toFixed(6)}</Text>
+            </HStack>
+          </VStack>
+          <VStack>
+            <Text fontWeight="bold">Tokens</Text>
+            <HStack>
+              <VStack w="28">
+                <Text>Prompt</Text>
+                <Text>{promptTokens ?? 0}</Text>
+              </VStack>
+              <Divider borderColor="gray.200" h={12} orientation="vertical" />
+              <VStack w="28">
+                <Text whiteSpace="nowrap">Completion</Text>
+                <Text>{completionTokens ?? 0}</Text>
+              </VStack>
+            </HStack>
+          </VStack>
+        </VStack>
+      }
+      {...props}
+    >
+      {children}
+    </Tooltip>
+  );
+};
diff --git a/src/server/api/routers/promptVariants.router.ts b/src/server/api/routers/promptVariants.router.ts
index 2e7db3e..1110d61 100644
--- a/src/server/api/routers/promptVariants.router.ts
+++ b/src/server/api/routers/promptVariants.router.ts
@@ -66,7 +66,7 @@ export const promptVariantsRouter = createTRPCRouter({
 
     const overallCost = overallPromptCost + overallCompletionCost;
 
-    return { evalResults, overallCost, scenarioCount, outputCount };
+    return { evalResults, promptTokens, completionTokens, overallCost, scenarioCount, outputCount };
   }),
 
   create: publicProcedure