Paginate scenarios

Show 10 scenarios at a time and let the user paginate through them to keep the interface responsive with potentially 1000s of scenarios.

.vscode/settings.json

@@ -1,6 +1,3 @@
 {
-  "eslint.format.enable": true,
-  "editor.codeActionsOnSave": {
-    "source.fixAll.eslint": true
-  }
+  "eslint.format.enable": true
 }

package.json

@@ -59,6 +59,7 @@
     "lodash-es": "^4.17.21",
     "next": "^13.4.2",
     "next-auth": "^4.22.1",
+    "next-query-params": "^4.2.3",
     "nextjs-routes": "^2.0.1",
     "openai": "4.0.0-beta.2",
     "pluralize": "^8.0.0",
@@ -79,6 +80,7 @@
     "superjson": "1.12.2",
     "tsx": "^3.12.7",
     "type-fest": "^4.0.0",
+    "use-query-params": "^2.2.1",
     "vite-tsconfig-paths": "^4.2.0",
     "zod": "^3.21.4",
     "zustand": "^4.3.9"

pnpm-lock.yaml

@@ -119,6 +119,9 @@ dependencies:
   next-auth:
     specifier: ^4.22.1
     version: 4.22.1(next@13.4.2)(react-dom@18.2.0)(react@18.2.0)
+  next-query-params:
+    specifier: ^4.2.3
+    version: 4.2.3(next@13.4.2)(react@18.2.0)(use-query-params@2.2.1)
   nextjs-routes:
     specifier: ^2.0.1
     version: 2.0.1(next@13.4.2)
@@ -179,6 +182,9 @@ dependencies:
   type-fest:
     specifier: ^4.0.0
     version: 4.0.0
+  use-query-params:
+    specifier: ^2.2.1
+    version: 2.2.1(react-dom@18.2.0)(react@18.2.0)
   vite-tsconfig-paths:
     specifier: ^4.2.0
     version: 4.2.0(typescript@5.0.4)
@@ -6037,6 +6043,19 @@ packages:
       uuid: 8.3.2
     dev: false
 
+  /next-query-params@4.2.3(next@13.4.2)(react@18.2.0)(use-query-params@2.2.1):
+    resolution: {integrity: sha512-hGNCYRH8YyA5ItiBGSKrtMl21b2MAqfPkdI1mvwloNVqSU142IaGzqHN+OTovyeLIpQfonY01y7BAHb/UH4POg==}
+    peerDependencies:
+      next: ^10.0.0 || ^11.0.0 || ^12.0.0 || ^13.0.0
+      react: ^16.8.0 || ^17.0.0 || ^18.0.0
+      use-query-params: ^2.0.0
+    dependencies:
+      next: 13.4.2(@babel/core@7.22.9)(react-dom@18.2.0)(react@18.2.0)
+      react: 18.2.0
+      tslib: 2.6.0
+      use-query-params: 2.2.1(react-dom@18.2.0)(react@18.2.0)
+    dev: false
+
   /next-tick@1.1.0:
     resolution: {integrity: sha512-CXdUiJembsNjuToQvxayPZF9Vqht7hewsvy2sOWafLvi2awflj9mOC6bHIg50orX8IJvWKY9wYQ/zB2kogPslQ==}
     dev: false
@@ -7147,6 +7166,10 @@ packages:
       randombytes: 2.1.0
     dev: true
 
+  /serialize-query-params@2.0.2:
+    resolution: {integrity: sha512-1chMo1dST4pFA9RDXAtF0Rbjaut4is7bzFbI1Z26IuMub68pNCILku85aYmeFhvnY//BXUPUhoRMjYcsT93J/Q==}
+    dev: false
+
   /serve-static@1.15.0:
     resolution: {integrity: sha512-XGuRDNjXUijsUL0vl6nSD7cwURuzEgglbOaFuZM9g3kwDXOWVTck0jLzjPzGD+TazWbboZYu52/9/XPdUgne9g==}
     engines: {node: '>= 0.8.0'}
@@ -7824,6 +7847,24 @@ packages:
       use-isomorphic-layout-effect: 1.1.2(@types/react@18.2.6)(react@18.2.0)
     dev: false
 
+  /use-query-params@2.2.1(react-dom@18.2.0)(react@18.2.0):
+    resolution: {integrity: sha512-i6alcyLB8w9i3ZK3caNftdb+UnbfBRNPDnc89CNQWkGRmDrm/gfydHvMBfVsQJRq3NoHOM2dt/ceBWG2397v1Q==}
+    peerDependencies:
+      '@reach/router': ^1.2.1
+      react: '>=16.8.0'
+      react-dom: '>=16.8.0'
+      react-router-dom: '>=5'
+    peerDependenciesMeta:
+      '@reach/router':
+        optional: true
+      react-router-dom:
+        optional: true
+    dependencies:
+      react: 18.2.0
+      react-dom: 18.2.0(react@18.2.0)
+      serialize-query-params: 2.0.2
+    dev: false
+
   /use-sidecar@1.1.2(@types/react@18.2.6)(react@18.2.0):
     resolution: {integrity: sha512-epTbsLuzZ7lPClpz2TyryBfztm7m+28DlEv2ZCQ3MDr5ssiwyOwGH/e5F9CkfWjJ1t4clvI58yF822/GUkjjhw==}
     engines: {node: '>=10'}

prisma/seed.ts

@@ -7,9 +7,13 @@ const defaultId = "11111111-1111-1111-1111-111111111111";
 await prisma.organization.deleteMany({
   where: { id: defaultId },
 });
-await prisma.organization.create({
-  data: { id: defaultId },
-});
+
+// If there's an existing org, just seed into it
+const org =
+  (await prisma.organization.findFirst({})) ??
+  (await prisma.organization.create({
+    data: { id: defaultId },
+  }));
 
 await prisma.experiment.deleteMany({
   where: {
@@ -21,7 +25,7 @@ await prisma.experiment.create({
   data: {
     id: defaultId,
     label: "Country Capitals Example",
-    organizationId: defaultId,
+    organizationId: org.id,
   },
 });
 
@@ -103,30 +107,41 @@ await prisma.testScenario.deleteMany({
   },
 });
 
+const countries = [
+  "Afghanistan",
+  "Albania",
+  "Algeria",
+  "Andorra",
+  "Angola",
+  "Antigua and Barbuda",
+  "Argentina",
+  "Armenia",
+  "Australia",
+  "Austria",
+  "Austrian Empire",
+  "Azerbaijan",
+  "Baden",
+  "Bahamas, The",
+  "Bahrain",
+  "Bangladesh",
+  "Barbados",
+  "Bavaria",
+  "Belarus",
+  "Belgium",
+  "Belize",
+  "Benin (Dahomey)",
+  "Bolivia",
+  "Bosnia and Herzegovina",
+  "Botswana",
+];
 await prisma.testScenario.createMany({
-  data: [
-    {
-      experimentId: defaultId,
-      sortIndex: 0,
-      variableValues: {
-        country: "Spain",
-      },
+  data: countries.map((country, i) => ({
+    experimentId: defaultId,
+    sortIndex: i,
+    variableValues: {
+      country: country,
     },
-    {
-      experimentId: defaultId,
-      sortIndex: 1,
-      variableValues: {
-        country: "USA",
-      },
-    },
-    {
-      experimentId: defaultId,
-      sortIndex: 2,
-      variableValues: {
-        country: "Chile",
-      },
-    },
-  ],
+  })),
 });
 
 const variants = await prisma.promptVariant.findMany({

src/components/OutputsTable/ScenarioPaginator.tsx

@@ -0,0 +1,74 @@
+import { Box, HStack, IconButton } from "@chakra-ui/react";
+import {
+  BsChevronDoubleLeft,
+  BsChevronDoubleRight,
+  BsChevronLeft,
+  BsChevronRight,
+} from "react-icons/bs";
+import { usePage, useScenarios } from "~/utils/hooks";
+
+const ScenarioPaginator = () => {
+  const [page, setPage] = usePage();
+  const { data } = useScenarios();
+
+  if (!data) return null;
+
+  const { scenarios, startIndex, lastPage, count } = data;
+
+  const nextPage = () => {
+    if (page < lastPage) {
+      setPage(page + 1, "replace");
+    }
+  };
+
+  const prevPage = () => {
+    if (page > 1) {
+      setPage(page - 1, "replace");
+    }
+  };
+
+  const goToLastPage = () => setPage(lastPage, "replace");
+  const goToFirstPage = () => setPage(1, "replace");
+
+  return (
+    <HStack pt={4}>
+      <IconButton
+        variant="ghost"
+        size="sm"
+        onClick={goToFirstPage}
+        isDisabled={page === 1}
+        aria-label="Go to first page"
+        icon={<BsChevronDoubleLeft />}
+      />
+      <IconButton
+        variant="ghost"
+        size="sm"
+        onClick={prevPage}
+        isDisabled={page === 1}
+        aria-label="Previous page"
+        icon={<BsChevronLeft />}
+      />
+      <Box>
+        {startIndex}-{startIndex + scenarios.length - 1} / {count}
+      </Box>
+      <IconButton
+        variant="ghost"
+        size="sm"
+        onClick={nextPage}
+        isDisabled={page === lastPage}
+        aria-label="Next page"
+        icon={<BsChevronRight />}
+      />
+      <IconButton
+        variant="ghost"
+        size="sm"
+        onClick={goToLastPage}
+        isDisabled={page === lastPage}
+        aria-label="Go to last page"
+        icon={<BsChevronDoubleRight />}
+      />
+    </HStack>
+  );
+};
+
+export default ScenarioPaginator;

src/components/OutputsTable/ScenariosHeader.tsx

@@ -12,7 +12,12 @@ import {
   Spinner,
 } from "@chakra-ui/react";
 import { cellPadding } from "../constants";
-import { useExperiment, useExperimentAccess, useHandledAsyncCallback } from "~/utils/hooks";
+import {
+  useExperiment,
+  useExperimentAccess,
+  useHandledAsyncCallback,
+  useScenarios,
+} from "~/utils/hooks";
 import { BsGear, BsPencil, BsPlus, BsStars } from "react-icons/bs";
 import { useAppStore } from "~/state/store";
 import { api } from "~/utils/api";
@@ -21,9 +26,10 @@ export const ActionButton = (props: ButtonProps) => (
   <Button size="sm" variant="ghost" color="gray.600" {...props} />
 );
 
-export const ScenariosHeader = (props: { numScenarios: number }) => {
+export const ScenariosHeader = () => {
   const openDrawer = useAppStore((s) => s.openDrawer);
   const { canModify } = useExperimentAccess();
+  const scenarios = useScenarios();
 
   const experiment = useExperiment();
   const createScenarioMutation = api.scenarios.create.useMutation();
@@ -44,7 +50,7 @@ export const ScenariosHeader = (props: { numScenarios: number }) => {
   return (
     <HStack w="100%" pb={cellPadding.y} pt={0} align="center" spacing={0}>
       <Text fontSize={16} fontWeight="bold">
-        Scenarios ({props.numScenarios})
+        Scenarios ({scenarios.data?.count})
       </Text>
       {canModify && (
         <Menu>

src/components/OutputsTable/index.tsx

@@ -7,6 +7,8 @@ import VariantHeader from "../VariantHeader/VariantHeader";
 import VariantStats from "./VariantStats";
 import { ScenariosHeader } from "./ScenariosHeader";
 import { borders } from "./styles";
+import { useScenarios } from "~/utils/hooks";
+import ScenarioPaginator from "./ScenarioPaginator";
 
 export default function OutputsTable({ experimentId }: { experimentId: string | undefined }) {
   const variants = api.promptVariants.list.useQuery(
@@ -14,17 +16,17 @@ export default function OutputsTable({ experimentId }: { experimentId: string |
     { enabled: !!experimentId },
   );
 
-  const scenarios = api.scenarios.list.useQuery(
-    { experimentId: experimentId as string },
-    { enabled: !!experimentId },
-  );
+  const scenarios = useScenarios();
 
   if (!variants.data || !scenarios.data) return null;
 
   const allCols = variants.data.length + 2;
   const variantHeaderRows = 3;
   const scenarioHeaderRows = 1;
-  const allRows = variantHeaderRows + scenarioHeaderRows + scenarios.data.length;
+  const scenarioFooterRows = 1;
+  const visibleScenariosCount = scenarios.data.scenarios.length;
+  const allRows =
+    variantHeaderRows + scenarioHeaderRows + visibleScenariosCount + scenarioFooterRows;
 
   return (
     <Grid
@@ -76,18 +78,25 @@ export default function OutputsTable({ experimentId }: { experimentId: string |
         {...borders}
         borderRightWidth={0}
       >
-        <ScenariosHeader numScenarios={scenarios.data.length} />
+        <ScenariosHeader />
       </GridItem>
 
-      {scenarios.data.map((scenario, i) => (
+      {scenarios.data.scenarios.map((scenario, i) => (
         <ScenarioRow
           rowStart={i + variantHeaderRows + scenarioHeaderRows + 2}
           key={scenario.uiId}
           scenario={scenario}
           variants={variants.data}
-          canHide={scenarios.data.length > 1}
+          canHide={visibleScenariosCount > 1}
         />
       ))}
+      <GridItem
+        rowStart={variantHeaderRows + scenarioHeaderRows + visibleScenariosCount + 2}
+        colStart={1}
+        colSpan={allCols}
+      >
+        <ScenarioPaginator />
+      </GridItem>
 
       {/* Add some extra padding on the right, because when the table is too wide to fit in the viewport `pr` on the Grid isn't respected. */}
       <GridItem rowStart={1} colStart={allCols} rowSpan={allRows} w={4} borderBottomWidth={0} />

src/components/OutputsTable/types.ts

@@ -2,4 +2,4 @@ import { type RouterOutputs } from "~/utils/api";
 
 export type PromptVariant = NonNullable<RouterOutputs["promptVariants"]["list"]>[0];
 
-export type Scenario = NonNullable<RouterOutputs["scenarios"]["list"]>[0];
+export type Scenario = NonNullable<RouterOutputs["scenarios"]["list"]>["scenarios"][0];

src/modelProviders/replicate-llama2/getCompletion.ts

@@ -27,8 +27,6 @@ export async function getCompletion(
       input: rest,
     });
 
-    console.log("stream?", onStream);
-
     const interval = onStream
       ? // eslint-disable-next-line @typescript-eslint/no-misused-promises
         setInterval(async () => {

src/pages/_app.tsx

@@ -7,6 +7,8 @@ import "~/utils/analytics";
 import Head from "next/head";
 import { ChakraThemeProvider } from "~/theme/ChakraThemeProvider";
 import { SyncAppStore } from "~/state/sync";
+import NextAdapterApp from "next-query-params/app";
+import { QueryParamProvider } from "use-query-params";
 
 const MyApp: AppType<{ session: Session | null }> = ({
   Component,
@@ -24,7 +26,9 @@ const MyApp: AppType<{ session: Session | null }> = ({
         <SyncAppStore />
         <Favicon />
         <ChakraThemeProvider>
-          <Component {...pageProps} />
+          <QueryParamProvider adapter={NextAdapterApp}>
+            <Component {...pageProps} />
+          </QueryParamProvider>
         </ChakraThemeProvider>
       </SessionProvider>
     </>

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

@@ -7,21 +7,39 @@ import { runAllEvals } from "~/server/utils/evaluations";
 import { generateNewCell } from "~/server/utils/generateNewCell";
 import { requireCanModifyExperiment, requireCanViewExperiment } from "~/utils/accessControl";
 
+const PAGE_SIZE = 10;
+
 export const scenariosRouter = createTRPCRouter({
   list: publicProcedure
-    .input(z.object({ experimentId: z.string() }))
+    .input(z.object({ experimentId: z.string(), page: z.number() }))
     .query(async ({ input, ctx }) => {
       await requireCanViewExperiment(input.experimentId, ctx);
 
-      return await prisma.testScenario.findMany({
+      const { experimentId, page } = input;
+
+      const scenarios = await prisma.testScenario.findMany({
         where: {
-          experimentId: input.experimentId,
+          experimentId,
           visible: true,
         },
-        orderBy: {
-          sortIndex: "asc",
+        orderBy: { sortIndex: "asc" },
+        skip: (page - 1) * PAGE_SIZE,
+        take: PAGE_SIZE,
+      });
+
+      const count = await prisma.testScenario.count({
+        where: {
+          experimentId,
+          visible: true,
         },
       });
+
+      return {
+        scenarios,
+        startIndex: (page - 1) * PAGE_SIZE + 1,
+        lastPage: Math.ceil(count / PAGE_SIZE),
+        count,
+      };
     }),
 
   create: protectedProcedure

src/state/sharedVariantEditor.slice.ts

@@ -8,9 +8,9 @@ export const editorBackground = "#fafafa";
 export type SharedVariantEditorSlice = {
   monaco: null | ReturnType<typeof loader.__getMonacoInstance>;
   loadMonaco: () => Promise<void>;
-  scenarios: RouterOutputs["scenarios"]["list"];
+  scenarios: RouterOutputs["scenarios"]["list"]["scenarios"];
   updateScenariosModel: () => void;
-  setScenarios: (scenarios: RouterOutputs["scenarios"]["list"]) => void;
+  setScenarios: (scenarios: RouterOutputs["scenarios"]["list"]["scenarios"]) => void;
 };
 
 export const createVariantEditorSlice: SliceCreator<SharedVariantEditorSlice> = (set, get) => ({

src/state/sync.tsx

@@ -1,17 +1,14 @@
 import { useEffect } from "react";
 import { api } from "~/utils/api";
-import { useExperiment } from "~/utils/hooks";
+import { useScenarios } from "~/utils/hooks";
 import { useAppStore } from "./store";
 
 export function useSyncVariantEditor() {
-  const experiment = useExperiment();
-  const scenarios = api.scenarios.list.useQuery(
-    { experimentId: experiment.data?.id ?? "" },
-    { enabled: !!experiment.data?.id },
-  );
+  const scenarios = useScenarios();
+
   useEffect(() => {
     if (scenarios.data) {
-      useAppStore.getState().sharedVariantEditor.setScenarios(scenarios.data);
+      useAppStore.getState().sharedVariantEditor.setScenarios(scenarios.data.scenarios);
     }
   }, [scenarios.data]);
 }

src/utils/hooks.ts

@@ -1,6 +1,7 @@
 import { useRouter } from "next/router";
 import { type RefObject, useCallback, useEffect, useRef, useState } from "react";
 import { api } from "~/utils/api";
+import { NumberParam, useQueryParam, withDefault } from "use-query-params";
 
 export const useExperiment = () => {
   const router = useRouter();
@@ -93,3 +94,15 @@ export const useElementDimensions = (): [RefObject<HTMLElement>, Dimensions | un
 
   return [ref, dimensions];
 };
+
+export const usePage = () => useQueryParam("page", withDefault(NumberParam, 1));
+
+export const useScenarios = () => {
+  const experiment = useExperiment();
+  const [page] = usePage();
+
+  return api.scenarios.list.useQuery(
+    { experimentId: experiment.data?.id ?? "", page },
+    { enabled: experiment.data?.id != null },
+  );
+};