Continue polling VariantStats while LLM retrieval in progress, minor UI fixes (#54)

* Prevent zoom in on iOS

* Expand function return code background to fill cell

* Keep OutputStats on far right of cells

* Continue polling prompt stats while cells are retrieving from LLM

* Add comment to _document.tsx

* Fix prettier

src/components/OutputsTable/OutputCell/OutputCell.tsx

@@ -1,6 +1,6 @@
 import { api } from "~/utils/api";
 import { type PromptVariant, type Scenario } from "../types";
-import { Spinner, Text, Box, Center, Flex, VStack } from "@chakra-ui/react";
+import { Spinner, Text, Center, VStack } from "@chakra-ui/react";
 import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
 import SyntaxHighlighter from "react-syntax-highlighter";
 import { docco } from "react-syntax-highlighter/dist/cjs/styles/hljs";
@@ -55,7 +55,10 @@ export default function OutputCell({
   const fetchingOutput = queryLoading || refetchingOutput;
 
   const awaitingOutput =
-    !cell || cell.retrievalStatus === "PENDING" || cell.retrievalStatus === "IN_PROGRESS";
+    !cell ||
+    cell.retrievalStatus === "PENDING" ||
+    cell.retrievalStatus === "IN_PROGRESS" ||
+    refetchingOutput;
   useEffect(() => setRefetchInterval(awaitingOutput ? 1000 : 0), [awaitingOutput]);
 
   const modelOutput = cell?.modelOutput;
@@ -95,11 +98,18 @@ export default function OutputCell({
     }
 
     return (
-      <Box fontSize="xs" width="100%" flexWrap="wrap" overflowX="auto">
+      <VStack
-        <VStack w="full" spacing={0}>
+        w="100%"
+        h="100%"
+        fontSize="xs"
+        flexWrap="wrap"
+        overflowX="auto"
+        justifyContent="space-between"
+      >
+        <VStack w="full" flex={1} spacing={0}>
           <CellOptions refetchingOutput={refetchingOutput} refetchOutput={hardRefetch} />
           <SyntaxHighlighter
-            customStyle={{ overflowX: "unset" }}
+            customStyle={{ overflowX: "unset", width: "100%", flex: 1 }}
             language="json"
             style={docco}
             lineProps={{
@@ -117,7 +127,7 @@ export default function OutputCell({
           </SyntaxHighlighter>
         </VStack>
         <OutputStats model={variant.model} modelOutput={modelOutput} scenario={scenario} />
-      </Box>
+      </VStack>
     );
   }
 
@@ -125,7 +135,7 @@ export default function OutputCell({
     message?.content ?? streamedContent ?? JSON.stringify(modelOutput?.output);
 
   return (
-    <Flex w="100%" h="100%" direction="column" justifyContent="space-between" whiteSpace="pre-wrap">
+    <VStack w="100%" h="100%" justifyContent="space-between" whiteSpace="pre-wrap">
       <VStack w="full" alignItems="flex-start" spacing={0}>
         <CellOptions refetchingOutput={refetchingOutput} refetchOutput={hardRefetch} />
         <Text>{contentToDisplay}</Text>
@@ -133,6 +143,6 @@ export default function OutputCell({
       {modelOutput && (
         <OutputStats model={variant.model} modelOutput={modelOutput} scenario={scenario} />
       )}
-    </Flex>
+    </VStack>
   );
 }

src/components/OutputsTable/OutputCell/OutputStats.tsx

@@ -31,7 +31,7 @@ export const OutputStats = ({
   const cost = promptCost + completionCost;
 
   return (
-    <HStack align="center" color="gray.500" fontSize="2xs" mt={{ base: 0, md: 1 }}>
+    <HStack w="full" align="center" color="gray.500" fontSize="2xs" mt={{ base: 0, md: 1 }}>
       <HStack flex={1}>
         {modelOutput.outputEvaluation.map((evaluation) => {
           const passed = evaluation.result > 0.5;

src/components/OutputsTable/VariantStats.tsx

@@ -1,12 +1,14 @@
-import { HStack, Icon, Text, useToken } from "@chakra-ui/react";
+import { HStack, Icon, Skeleton, Text, useToken } from "@chakra-ui/react";
 import { type PromptVariant } from "./types";
 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";
+import { useEffect, useState } from "react";
 
 export default function VariantStats(props: { variant: PromptVariant }) {
+  const [refetchInterval, setRefetchInterval] = useState(0);
   const { data } = api.promptVariants.stats.useQuery(
     {
       variantId: props.variant.id,
@@ -19,10 +21,18 @@ export default function VariantStats(props: { variant: PromptVariant }) {
         completionTokens: 0,
         scenarioCount: 0,
         outputCount: 0,
+        awaitingRetrievals: false,
       },
+      refetchInterval,
     },
   );
 
+  // Poll every two seconds while we are waiting for LLM retrievals to finish
+  useEffect(
+    () => setRefetchInterval(data.awaitingRetrievals ? 2000 : 0),
+    [data.awaitingRetrievals],
+  );
+
   const [passColor, neutralColor, failColor] = useToken("colors", [
     "green.500",
     "gray.500",
@@ -33,16 +43,20 @@ export default function VariantStats(props: { variant: PromptVariant }) {
 
   const showNumFinished = data.scenarioCount > 0 && data.scenarioCount !== data.outputCount;
 
-  if (!(data.evalResults.length > 0) && !data.overallCost) return null;
-
   return (
-    <HStack justifyContent="space-between" alignItems="center" mx="2" fontSize="xs">
+    <HStack
+      justifyContent="space-between"
+      alignItems="center"
+      mx="2"
+      fontSize="xs"
+      py={cellPadding.y}
+    >
       {showNumFinished && (
         <Text>
           {data.outputCount} / {data.scenarioCount}
         </Text>
       )}
-      <HStack px={cellPadding.x} py={cellPadding.y}>
+      <HStack px={cellPadding.x}>
         {data.evalResults.map((result) => {
           const passedFrac = result.passCount / result.totalCount;
           return (
@@ -55,17 +69,19 @@ export default function VariantStats(props: { variant: PromptVariant }) {
           );
         })}
       </HStack>
-      {data.overallCost && (
+      {data.overallCost && !data.awaitingRetrievals ? (
         <CostTooltip
           promptTokens={data.promptTokens}
           completionTokens={data.completionTokens}
           cost={data.overallCost}
         >
-          <HStack spacing={0} align="center" color="gray.500" my="2">
+          <HStack spacing={0} align="center" color="gray.500">
             <Icon as={BsCurrencyDollar} />
             <Text mr={1}>{data.overallCost.toFixed(3)}</Text>
           </HStack>
         </CostTooltip>
+      ) : (
+        <Skeleton height={4} width={12} mr={1} />
       )}
     </HStack>
   );

src/components/tooltip/CostTooltip.tsx

@@ -20,7 +20,6 @@ export const CostTooltip = ({
       color="gray.800"
       bgColor="gray.50"
       borderWidth={1}
-      py={2}
       hasArrow
       shouldWrapChildren
       label={

src/pages/_document.tsx

@@ -0,0 +1,23 @@
+import Document, { Html, Head, Main, NextScript } from "next/document";
+
+class MyDocument extends Document {
+  render() {
+    return (
+      <Html>
+        <Head>
+          {/* Prevent automatic zoom-in on iPhone when focusing on text input */}
+          <meta
+            name="viewport"
+            content="width=device-width, initial-scale=1, maximum-scale=1, user-scalable=0"
+          />
+        </Head>
+        <body>
+          <Main />
+          <NextScript />
+        </body>
+      </Html>
+    );
+  }
+}
+
+export default MyDocument;

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

@@ -109,7 +109,26 @@ export const promptVariantsRouter = createTRPCRouter({
 
     const overallCost = overallPromptCost + overallCompletionCost;
 
-    return { evalResults, promptTokens, completionTokens, overallCost, scenarioCount, outputCount };
+    const awaitingRetrievals = !!(await prisma.scenarioVariantCell.findFirst({
+      where: {
+        promptVariantId: input.variantId,
+        testScenario: { visible: true },
+        // Check if is PENDING or IN_PROGRESS
+        retrievalStatus: {
+          in: ["PENDING", "IN_PROGRESS"],
+        },
+      },
+    }));
+
+    return {
+      evalResults,
+      promptTokens,
+      completionTokens,
+      overallCost,
+      scenarioCount,
+      outputCount,
+      awaitingRetrievals,
+    };
   }),
 
   create: publicProcedure