Merge pull request #3421 from quantified-uncertainty/ai-pagination

/ai pagination

packages/hub/src/app/ai/AiDashboard.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { clsx } from "clsx";
+import clsx from "clsx";
 import { FC, useRef, useState } from "react";
 
 import { ClientWorkflow, ClientWorkflowResult } from "@quri/squiggle-ai";
@@ -14,9 +14,15 @@ export type SquiggleResponse = {
   currentStep?: string;
 };
 
-export const AiDashboard: FC<{ initialWorkflows: ClientWorkflow[] }> = ({
+type Props = {
+  initialWorkflows: ClientWorkflow[];
+  hasMoreWorkflows: boolean;
+};
+
+export const AiDashboard: FC<Props> = ({
   initialWorkflows,
-}) => {
+  hasMoreWorkflows,
+}: Props) => {
   const { workflows, submitWorkflow, selectedWorkflow, selectWorkflow } =
     useSquiggleWorkflows(initialWorkflows);
 
@@ -38,6 +44,7 @@ export const AiDashboard: FC<{ initialWorkflows: ClientWorkflow[] }> = ({
           selectWorkflow={selectWorkflow}
           selectedWorkflow={selectedWorkflow}
           workflows={workflows}
+          hasMoreWorkflows={hasMoreWorkflows}
           ref={sidebarRef}
         />
       </div>

packages/hub/src/app/ai/Sidebar.tsx

@@ -17,6 +17,8 @@ import {
   TextAreaFormField,
 } from "@quri/ui";
 
+import { LoadMoreViaSearchParam } from "@/components/LoadMoreViaSearchParam";
+
 import { AiRequestBody } from "./utils";
 import { WorkflowSummaryList } from "./WorkflowSummaryList";
 
@@ -29,6 +31,7 @@ type Props = {
   selectWorkflow: (id: string) => void;
   selectedWorkflow: ClientWorkflow | undefined;
   workflows: ClientWorkflow[];
+  hasMoreWorkflows: boolean;
 };
 
 type FormShape = {
@@ -38,7 +41,13 @@ type FormShape = {
 };
 
 export const Sidebar = forwardRef<Handle, Props>(function Sidebar(
-  { submitWorkflow, selectWorkflow, selectedWorkflow, workflows },
+  {
+    submitWorkflow,
+    selectWorkflow,
+    selectedWorkflow,
+    workflows,
+    hasMoreWorkflows,
+  },
   ref
 ) {
   const form = useForm<FormShape>({
@@ -162,6 +171,7 @@ Outputs:
             selectedWorkflow={selectedWorkflow}
             selectWorkflow={selectWorkflow}
           />
+          {hasMoreWorkflows && <LoadMoreViaSearchParam />}
         </div>
       </div>
     </FormProvider>

packages/hub/src/app/ai/WorkflowSummaryList.tsx

@@ -1,4 +1,3 @@
-import { orderBy } from "lodash";
 import { FC } from "react";
 
 import { ClientWorkflow } from "@quri/squiggle-ai";
@@ -10,11 +9,9 @@ export const WorkflowSummaryList: FC<{
   selectedWorkflow: ClientWorkflow | undefined;
   selectWorkflow: (id: string) => void;
 }> = ({ workflows, selectedWorkflow, selectWorkflow }) => {
-  const sortedWorkflows = orderBy(workflows, ["timestamp"], ["desc"]);
-
   return (
     <div className="flex max-h-[400px] flex-col space-y-2 overflow-y-auto pr-2">
-      {sortedWorkflows.map((workflow) => (
+      {workflows.map((workflow) => (
         <WorkflowSummaryItem
           key={workflow.id}
           workflow={workflow}

packages/hub/src/app/ai/page.tsx

@@ -1,24 +1,24 @@
-import { ClientWorkflow } from "@quri/squiggle-ai";
+import { z } from "zod";
 
-import { prisma } from "@/prisma";
-import { getUserOrRedirect } from "@/server/helpers";
+import { numberInString } from "@/lib/zodUtils";
+import { loadWorkflows } from "@/server/ai/data";
 
-import { decodeDbWorkflowToClientWorkflow } from "../../server/ai/storage";
 import { AiDashboard } from "./AiDashboard";
 
-export default async function AiPage() {
-  const user = await getUserOrRedirect();
+export default async function SessionsPage({
+  searchParams,
+}: {
+  searchParams: { [key: string]: string | string[] | undefined };
+}) {
+  const { limit } = z
+    .object({
+      limit: numberInString.optional(),
+    })
+    .parse(searchParams);
 
-  const rows = await prisma.aiWorkflow.findMany({
-    orderBy: { createdAt: "desc" },
-    where: {
-      user: { email: user.email },
-    },
-  });
+  const { workflows, hasMore } = await loadWorkflows({ limit });
 
-  const workflows: ClientWorkflow[] = rows.map((row) =>
-    decodeDbWorkflowToClientWorkflow(row)
+  return (
+    <AiDashboard initialWorkflows={workflows} hasMoreWorkflows={hasMore} />
   );
-
-  return <AiDashboard initialWorkflows={workflows} />;
 }

packages/hub/src/app/ai/useSquiggleWorkflows.tsx

@@ -1,14 +1,26 @@
-import { useCallback, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 
 import { ClientWorkflow, decodeWorkflowFromReader } from "@quri/squiggle-ai";
 
 import { AiRequestBody, bodyToLineReader } from "./utils";
 
-export function useSquiggleWorkflows(initialWorkflows: ClientWorkflow[]) {
+export function useSquiggleWorkflows(preloadedWorkflows: ClientWorkflow[]) {
   const [workflows, setWorkflows] =
-    useState<ClientWorkflow[]>(initialWorkflows);
+    useState<ClientWorkflow[]>(preloadedWorkflows);
   const [selected, setSelected] = useState<number | undefined>(undefined);
 
+  // `preloadedWorkflows` can change when the user presses the "load more" button
+  useEffect(() => {
+    setWorkflows((list) => {
+      if (list === preloadedWorkflows) return list;
+      const knownWorkflows = new Set(list.map((w) => w.id));
+      const newWorkflows = preloadedWorkflows.filter(
+        (w) => !knownWorkflows.has(w.id)
+      );
+      return [...list, ...newWorkflows];
+    });
+  }, [preloadedWorkflows]);
+
   const updateWorkflow = useCallback(
     (id: string, update: (workflow: ClientWorkflow) => ClientWorkflow) => {
       setWorkflows((workflows) =>
@@ -20,29 +32,26 @@ export function useSquiggleWorkflows(initialWorkflows: ClientWorkflow[]) {
     []
   );
 
-  const addMockWorkflow = useCallback(
-    (request: AiRequestBody) => {
-      // This will be replaced with a real workflow once we receive the first message from the server.
-      const id = `loading-${Date.now().toString()}`;
-      const workflow: ClientWorkflow = {
-        id,
-        timestamp: new Date().getTime(),
-        status: "loading",
-        inputs: {
-          prompt: {
-            id: "prompt",
-            kind: "prompt",
-            value: request.kind === "create" ? request.prompt : "[FIX]",
-          },
+  const addMockWorkflow = useCallback((request: AiRequestBody) => {
+    // This will be replaced with a real workflow once we receive the first message from the server.
+    const id = `loading-${Date.now().toString()}`;
+    const workflow: ClientWorkflow = {
+      id,
+      timestamp: new Date().getTime(),
+      status: "loading",
+      inputs: {
+        prompt: {
+          id: "prompt",
+          kind: "prompt",
+          value: request.kind === "create" ? request.prompt : "[FIX]",
         },
-        steps: [],
-      };
-      setWorkflows((workflows) => [...workflows, workflow]);
-      setSelected(workflows.length); // select the new workflow
-      return workflow;
-    },
-    [workflows.length]
-  );
+      },
+      steps: [],
+    };
+    setWorkflows((workflows) => [workflow, ...workflows]);
+    setSelected(0);
+    return workflow;
+  }, []);
 
   const submitWorkflow = useCallback(
     async (request: AiRequestBody) => {

packages/hub/src/components/LoadMoreViaSearchParam.tsx

@@ -0,0 +1,29 @@
+import { FC } from "react";
+
+import { useUpdateSearchParams } from "@/hooks/useUpdateSearchParams";
+
+import { LoadMore } from "./LoadMore";
+
+type Props = {
+  param?: string;
+};
+
+export const LoadMoreViaSearchParam: FC<Props> = ({ param = "limit" }) => {
+  const updateSearchParams = useUpdateSearchParams();
+
+  const action = (count: number) => {
+    updateSearchParams(
+      (params) => {
+        if (params.get(param)) {
+          const oldValue = parseInt(params.get(param) as string);
+          params.set(param, String(oldValue + count));
+        } else {
+          params.set(param, String(count));
+        }
+      },
+      { mode: "replace", scroll: false }
+    );
+  };
+
+  return <LoadMore loadNext={action} />;
+};

packages/hub/src/hooks/useUpdateSearchParams.ts

@@ -0,0 +1,29 @@
+import { usePathname, useRouter, useSearchParams } from "next/navigation";
+import { useCallback } from "react";
+
+export function useUpdateSearchParams() {
+  const router = useRouter();
+  const pathname = usePathname();
+  const searchParams = useSearchParams();
+
+  const updateSearchParams = useCallback(
+    (
+      update: (params: URLSearchParams) => void,
+      {
+        mode = "push",
+        scroll = true,
+      }: {
+        mode?: "push" | "replace";
+        scroll?: boolean;
+      } = {}
+    ) => {
+      const currentParams = new URLSearchParams(searchParams.toString());
+      update(currentParams);
+      const method = mode === "push" ? router.push : router.replace;
+      method(`${pathname}?${currentParams}`, { scroll });
+    },
+    [pathname, searchParams, router]
+  );
+
+  return updateSearchParams;
+}

packages/hub/src/lib/zodUtils.ts

@@ -0,0 +1,18 @@
+import { z } from "zod";
+
+export const numberInString = z.string().transform((val, ctx) => {
+  const parsed = parseInt(val);
+  if (isNaN(parsed)) {
+    ctx.addIssue({
+      code: z.ZodIssueCode.custom,
+      message: "Not a number",
+    });
+
+    // This is a special symbol you can use to
+    // return early from the transform function.
+    // It has type `never` so it does not affect the
+    // inferred return type.
+    return z.NEVER;
+  }
+  return parsed;
+});

packages/hub/src/server/ai/data.ts

@@ -0,0 +1,29 @@
+import "server-only";
+
+import { prisma } from "@/prisma";
+
+import { getUserOrRedirect } from "../helpers";
+import { decodeDbWorkflowToClientWorkflow } from "./storage";
+
+export async function loadWorkflows({
+  limit = 20,
+}: {
+  limit?: number;
+} = {}) {
+  const user = await getUserOrRedirect();
+
+  const rows = await prisma.aiWorkflow.findMany({
+    orderBy: { createdAt: "desc" },
+    where: {
+      user: { email: user.email },
+    },
+    take: limit + 1,
+  });
+
+  const workflows = rows.map((row) => decodeDbWorkflowToClientWorkflow(row));
+
+  return {
+    workflows: limit ? workflows.slice(0, limit) : workflows,
+    hasMore: limit ? workflows.length > limit : false,
+  };
+}