Allow user to toggle visible columns (#182)

* Maintain tag casing

* Persist column visibility in zustand

* Persist only visibleColumns key

* merge persisted state

* Only show ColumnVisibilityDropdown after rehydration

* Record storage rehydrated

* Add useIsClientRehydrated hook

* Hide ActionButton text on mobile

* Condense Paginator on mobile

---------

Co-authored-by: Kyle Corbitt <kyle@corbt.com>

app/src/components/Paginator.tsx

@@ -1,15 +1,19 @@
-import { HStack, IconButton, Text, Select, type StackProps, Icon } from "@chakra-ui/react";
+import {
+  HStack,
+  IconButton,
+  Text,
+  Select,
+  type StackProps,
+  Icon,
+  useBreakpointValue,
+} from "@chakra-ui/react";
 import React, { useCallback } from "react";
 import { FiChevronsLeft, FiChevronsRight, FiChevronLeft, FiChevronRight } from "react-icons/fi";
 import { usePageParams } from "~/utils/hooks";
 
 const pageSizeOptions = [10, 25, 50, 100];
 
-const Paginator = ({
+const Paginator = ({ count, ...props }: { count: number; condense?: boolean } & StackProps) => {
-  count,
-  condense,
-  ...props
-}: { count: number; condense?: boolean } & StackProps) => {
   const { page, pageSize, setPageParams } = usePageParams();
 
   const lastPage = Math.ceil(count / pageSize);
@@ -37,6 +41,9 @@ const Paginator = ({
   const goToLastPage = () => setPageParams({ page: lastPage }, "replace");
   const goToFirstPage = () => setPageParams({ page: 1 }, "replace");
 
+  const isMobile = useBreakpointValue({ base: true, md: false });
+  const condense = isMobile || props.condense;
+
   if (count === 0) return null;
 
   return (

app/src/components/nav/UserMenu.tsx

@@ -23,7 +23,6 @@ export default function UserMenu({ user, ...rest }: { user: Session } & StackPro
   );
 
   return (
-    <>
     <Popover placement="right">
       <PopoverTrigger>
         <NavSidebarOption>
@@ -67,6 +66,5 @@ export default function UserMenu({ user, ...rest }: { user: Session } & StackPro
         </VStack>
       </PopoverContent>
     </Popover>
-    </>
   );
 }

app/src/components/requestLogs/ActionButton.tsx

@@ -21,7 +21,7 @@ const ActionButton = ({
     >
       <HStack spacing={1}>
         {icon && <Icon as={icon} />}
-        <Text>{label}</Text>
+        <Text display={{ base: "none", md: "flex" }}>{label}</Text>
       </HStack>
     </Button>
   );

app/src/components/requestLogs/ColumnVisiblityDropdown.tsx

@@ -0,0 +1,117 @@
+import {
+  Icon,
+  Popover,
+  PopoverTrigger,
+  PopoverContent,
+  VStack,
+  HStack,
+  Button,
+  Text,
+  useDisclosure,
+  Box,
+} from "@chakra-ui/react";
+import { BiCheck } from "react-icons/bi";
+import { BsToggles } from "react-icons/bs";
+import { useMemo } from "react";
+
+import { useIsClientRehydrated, useTagNames } from "~/utils/hooks";
+import { useAppStore } from "~/state/store";
+import { StaticColumnKeys } from "~/state/columnVisiblitySlice";
+import ActionButton from "./ActionButton";
+
+const ColumnVisiblityDropdown = () => {
+  const tagNames = useTagNames().data;
+
+  const visibleColumns = useAppStore((s) => s.columnVisibility.visibleColumns);
+  const toggleColumnVisibility = useAppStore((s) => s.columnVisibility.toggleColumnVisibility);
+  const totalColumns = Object.keys(StaticColumnKeys).length + (tagNames?.length ?? 0);
+
+  const popover = useDisclosure();
+
+  const columnVisiblityOptions = useMemo(() => {
+    const options: { label: string; key: string }[] = [
+      {
+        label: "Sent At",
+        key: StaticColumnKeys.SENT_AT,
+      },
+      {
+        label: "Model",
+        key: StaticColumnKeys.MODEL,
+      },
+      {
+        label: "Duration",
+        key: StaticColumnKeys.DURATION,
+      },
+      {
+        label: "Input Tokens",
+        key: StaticColumnKeys.INPUT_TOKENS,
+      },
+      {
+        label: "Output Tokens",
+        key: StaticColumnKeys.OUTPUT_TOKENS,
+      },
+      {
+        label: "Status Code",
+        key: StaticColumnKeys.STATUS_CODE,
+      },
+    ];
+    for (const tagName of tagNames ?? []) {
+      options.push({
+        label: tagName,
+        key: tagName,
+      });
+    }
+    return options;
+  }, [tagNames]);
+
+  const isClientRehydrated = useIsClientRehydrated();
+  if (!isClientRehydrated) return null;
+
+  return (
+    <Popover
+      placement="bottom-start"
+      isOpen={popover.isOpen}
+      onOpen={popover.onOpen}
+      onClose={popover.onClose}
+    >
+      <PopoverTrigger>
+        <Box>
+          <ActionButton
+            label={`Columns (${visibleColumns.size}/${totalColumns})`}
+            icon={BsToggles}
+          />
+        </Box>
+      </PopoverTrigger>
+      <PopoverContent boxShadow="0 0 40px 4px rgba(0, 0, 0, 0.1);" minW={0} w="auto">
+        <VStack spacing={0} maxH={400} overflowY="auto">
+          {columnVisiblityOptions?.map((option, index) => (
+            <HStack
+              key={index}
+              as={Button}
+              onClick={() => toggleColumnVisibility(option.key)}
+              w="full"
+              minH={10}
+              variant="ghost"
+              justifyContent="space-between"
+              fontWeight="semibold"
+              borderRadius={0}
+              colorScheme="blue"
+              color="black"
+              fontSize="sm"
+              borderBottomWidth={1}
+            >
+              <Text mr={16}>{option.label}</Text>
+              <Box w={5}>
+                {visibleColumns.has(option.key) && (
+                  <Icon as={BiCheck} color="blue.500" boxSize={5} />
+                )}
+              </Box>
+            </HStack>
+          ))}
+        </VStack>
+      </PopoverContent>
+    </Popover>
+  );
+};
+
+export default ColumnVisiblityDropdown;

app/src/components/requestLogs/LoggedCallsTable.tsx

@@ -10,7 +10,7 @@ export default function LoggedCallsTable() {
   return (
     <Card width="100%" overflowX="auto">
       <Table>
-        <TableHeader showCheckbox />
+        <TableHeader isSimple />
         <Tbody>
           {loggedCalls?.calls?.map((loggedCall) => {
             return (
@@ -25,7 +25,7 @@ export default function LoggedCallsTable() {
                     setExpandedRow(loggedCall.id);
                   }
                 }}
-                showCheckbox
+                isSimple
               />
             );
           })}

app/src/components/requestLogs/TableRow.tsx

@@ -21,14 +21,15 @@ import Link from "next/link";
 import { type RouterOutputs } from "~/utils/api";
 import { FormattedJson } from "./FormattedJson";
 import { useAppStore } from "~/state/store";
-import { useLoggedCalls, useTagNames } from "~/utils/hooks";
+import { useIsClientRehydrated, useLoggedCalls, useTagNames } from "~/utils/hooks";
 import { useMemo } from "react";
+import { StaticColumnKeys } from "~/state/columnVisiblitySlice";
 
 dayjs.extend(relativeTime);
 
 type LoggedCall = RouterOutputs["loggedCalls"]["list"]["calls"][0];
 
-export const TableHeader = ({ showCheckbox }: { showCheckbox?: boolean }) => {
+export const TableHeader = ({ isSimple }: { isSimple?: boolean }) => {
   const matchingLogIds = useLoggedCalls().data?.matchingLogIds;
   const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
   const addAll = useAppStore((s) => s.selectedLogs.addSelectedLogIds);
@@ -38,10 +39,14 @@ export const TableHeader = ({ showCheckbox }: { showCheckbox?: boolean }) => {
     return matchingLogIds.every((id) => selectedLogIds.has(id));
   }, [selectedLogIds, matchingLogIds]);
   const tagNames = useTagNames().data;
+  const visibleColumns = useAppStore((s) => s.columnVisibility.visibleColumns);
+  const isClientRehydrated = useIsClientRehydrated();
+  if (!isClientRehydrated) return null;
+
   return (
     <Thead>
       <Tr>
-        {showCheckbox && (
+        {isSimple && (
           <Th pr={0}>
             <HStack minW={16}>
               <Checkbox
@@ -57,13 +62,19 @@ export const TableHeader = ({ showCheckbox }: { showCheckbox?: boolean }) => {
             </HStack>
           </Th>
         )}
-        <Th>Sent At</Th>
+        {visibleColumns.has(StaticColumnKeys.SENT_AT) && <Th>Sent At</Th>}
-        <Th>Model</Th>
+        {visibleColumns.has(StaticColumnKeys.MODEL) && <Th>Model</Th>}
-        {tagNames?.map((tagName) => <Th key={tagName}>{tagName}</Th>)}
+        {tagNames
-        <Th isNumeric>Duration</Th>
+          ?.filter((tagName) => visibleColumns.has(tagName))
-        <Th isNumeric>Input tokens</Th>
+          .map((tagName) => (
-        <Th isNumeric>Output tokens</Th>
+            <Th key={tagName} textTransform={"none"}>
-        <Th isNumeric>Status</Th>
+              {tagName}
+            </Th>
+          ))}
+        {visibleColumns.has(StaticColumnKeys.DURATION) && <Th isNumeric>Duration</Th>}
+        {visibleColumns.has(StaticColumnKeys.INPUT_TOKENS) && <Th isNumeric>Input tokens</Th>}
+        {visibleColumns.has(StaticColumnKeys.OUTPUT_TOKENS) && <Th isNumeric>Output tokens</Th>}
+        {visibleColumns.has(StaticColumnKeys.STATUS_CODE) && <Th isNumeric>Status</Th>}
       </Tr>
     </Thead>
   );
@@ -73,12 +84,12 @@ export const TableRow = ({
   loggedCall,
   isExpanded,
   onToggle,
-  showCheckbox,
+  isSimple,
 }: {
   loggedCall: LoggedCall;
   isExpanded: boolean;
   onToggle: () => void;
-  showCheckbox?: boolean;
+  isSimple?: boolean;
 }) => {
   const isError = loggedCall.modelResponse?.statusCode !== 200;
   const requestedAt = dayjs(loggedCall.requestedAt).format("MMMM D h:mm A");
@@ -88,6 +99,10 @@ export const TableRow = ({
   const toggleChecked = useAppStore((s) => s.selectedLogs.toggleSelectedLogId);
 
   const tagNames = useTagNames().data;
+  const visibleColumns = useAppStore((s) => s.columnVisibility.visibleColumns);
+
+  const isClientRehydrated = useIsClientRehydrated();
+  if (!isClientRehydrated) return null;
 
   return (
     <>
@@ -100,11 +115,12 @@ export const TableRow = ({
         }}
         fontSize="sm"
       >
-        {showCheckbox && (
+        {isSimple && (
           <Td>
             <Checkbox isChecked={isChecked} onChange={() => toggleChecked(loggedCall.id)} />
           </Td>
         )}
+        {visibleColumns.has(StaticColumnKeys.SENT_AT) && (
           <Td>
             <Tooltip label={fullTime} placement="top">
               <Box whiteSpace="nowrap" minW="120px">
@@ -112,6 +128,8 @@ export const TableRow = ({
               </Box>
             </Tooltip>
           </Td>
+        )}
+        {visibleColumns.has(StaticColumnKeys.MODEL) && (
           <Td>
             <HStack justifyContent="flex-start">
               <Text
@@ -128,7 +146,11 @@ export const TableRow = ({
               </Text>
             </HStack>
           </Td>
-        {tagNames?.map((tagName) => <Td key={tagName}>{loggedCall.tags[tagName]}</Td>)}
+        )}
+        {tagNames
+          ?.filter((tagName) => visibleColumns.has(tagName))
+          .map((tagName) => <Td key={tagName}>{loggedCall.tags[tagName]}</Td>)}
+        {visibleColumns.has(StaticColumnKeys.DURATION) && (
           <Td isNumeric>
             {loggedCall.cacheHit ? (
               <Text color="gray.500">Cached</Text>
@@ -136,11 +158,18 @@ export const TableRow = ({
               ((loggedCall.modelResponse?.durationMs ?? 0) / 1000).toFixed(2) + "s"
             )}
           </Td>
+        )}
+        {visibleColumns.has(StaticColumnKeys.INPUT_TOKENS) && (
           <Td isNumeric>{loggedCall.modelResponse?.inputTokens}</Td>
+        )}
+        {visibleColumns.has(StaticColumnKeys.OUTPUT_TOKENS) && (
           <Td isNumeric>{loggedCall.modelResponse?.outputTokens}</Td>
+        )}
+        {visibleColumns.has(StaticColumnKeys.STATUS_CODE) && (
           <Td sx={{ color: isError ? "red.500" : "green.500", fontWeight: "semibold" }} isNumeric>
             {loggedCall.modelResponse?.statusCode ?? "No response"}
           </Td>
+        )}
       </Tr>
       <Tr>
         <Td colSpan={8} p={0}>

app/src/pages/request-logs/index.tsx

@@ -9,6 +9,7 @@ import { useAppStore } from "~/state/store";
 import { RiFlaskLine } from "react-icons/ri";
 import { FiFilter } from "react-icons/fi";
 import LogFilters from "~/components/requestLogs/LogFilters/LogFilters";
+import ColumnVisiblityDropdown from "~/components/requestLogs/ColumnVisiblityDropdown";
 
 export default function LoggedCalls() {
   const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
@@ -23,6 +24,7 @@ export default function LoggedCalls() {
         </Text>
         <Divider />
         <HStack w="full" justifyContent="flex-end">
+          <ColumnVisiblityDropdown />
           <ActionButton
             onClick={() => {
               setFiltersShown(!filtersShown);

app/src/state/columnVisiblitySlice.ts

@@ -0,0 +1,37 @@
+import { type SliceCreator } from "./store";
+
+export const comparators = ["=", "!=", "CONTAINS", "NOT_CONTAINS"] as const;
+
+export const defaultFilterableFields = ["Request", "Response", "Model", "Status Code"] as const;
+
+export enum StaticColumnKeys {
+  SENT_AT = "sentAt",
+  MODEL = "model",
+  DURATION = "duration",
+  INPUT_TOKENS = "inputTokens",
+  OUTPUT_TOKENS = "outputTokens",
+  STATUS_CODE = "statusCode",
+}
+
+export type ColumnVisibilitySlice = {
+  visibleColumns: Set<string>;
+  toggleColumnVisibility: (columnKey: string) => void;
+  showAllColumns: (columnKeys: string[]) => void;
+};
+
+export const createColumnVisibilitySlice: SliceCreator<ColumnVisibilitySlice> = (set, get) => ({
+  // initialize with all static columns visible
+  visibleColumns: new Set(Object.values(StaticColumnKeys)),
+  toggleColumnVisibility: (columnKey: string) =>
+    set((state) => {
+      if (state.columnVisibility.visibleColumns.has(columnKey)) {
+        state.columnVisibility.visibleColumns.delete(columnKey);
+      } else {
+        state.columnVisibility.visibleColumns.add(columnKey);
+      }
+    }),
+  showAllColumns: (columnKeys: string[]) =>
+    set((state) => {
+      state.columnVisibility.visibleColumns = new Set(columnKeys);
+    }),
+});

app/src/state/persist.ts

@@ -1,13 +1,27 @@
 import { type PersistOptions } from "zustand/middleware/persist";
 import { type State } from "./store";
+import SuperJSON from "superjson";
+import { merge, pick } from "lodash-es";
+import { type PartialDeep } from "type-fest";
 
-export const stateToPersist = {
+export type PersistedState = PartialDeep<State>;
-  selectedProjectId: null as string | null,
-};
 
-export const persistOptions: PersistOptions<State, typeof stateToPersist> = {
+export const persistOptions: PersistOptions<State, PersistedState> = {
   name: "persisted-app-store",
   partialize: (state) => ({
     selectedProjectId: state.selectedProjectId,
+    columnVisibility: pick(state.columnVisibility, ["visibleColumns"]),
   }),
+  merge: (saved, state) => merge(state, saved),
+  storage: {
+    getItem: (key) => {
+      const data = localStorage.getItem(key);
+      return data ? SuperJSON.parse(data) : null;
+    },
+    setItem: (key, value) => localStorage.setItem(key, SuperJSON.stringify(value)),
+    removeItem: (key) => localStorage.removeItem(key),
+  },
+  onRehydrateStorage: (state) => {
+    if (state) state.isRehydrated = true;
+  },
 };

app/src/state/store.ts

@@ -8,13 +8,15 @@ import {
   createVariantEditorSlice,
 } from "./sharedVariantEditor.slice";
 import { type APIClient } from "~/utils/api";
-import { persistOptions, type stateToPersist } from "./persist";
+import { type PersistedState, persistOptions } from "./persist";
 import { type SelectedLogsSlice, createSelectedLogsSlice } from "./selectedLogsSlice";
 import { type LogFiltersSlice, createLogFiltersSlice } from "./logFiltersSlice";
+import { createColumnVisibilitySlice, type ColumnVisibilitySlice } from "./columnVisiblitySlice";
 
 enableMapSet();
 
 export type State = {
+  isRehydrated: boolean;
   drawerOpen: boolean;
   openDrawer: () => void;
   closeDrawer: () => void;
@@ -25,6 +27,7 @@ export type State = {
   setSelectedProjectId: (id: string) => void;
   selectedLogs: SelectedLogsSlice;
   logFilters: LogFiltersSlice;
+  columnVisibility: ColumnVisibilitySlice;
 };
 
 export type SliceCreator<T> = StateCreator<State, [["zustand/immer", never]], [], T>;
@@ -32,18 +35,15 @@ export type SliceCreator<T> = StateCreator<State, [["zustand/immer", never]], []
 export type SetFn = Parameters<SliceCreator<unknown>>[0];
 export type GetFn = Parameters<SliceCreator<unknown>>[1];
 
-const useBaseStore = create<
+const useBaseStore = create<State, [["zustand/persist", PersistedState], ["zustand/immer", never]]>(
-  State,
-  [["zustand/persist", typeof stateToPersist], ["zustand/immer", never]]
->(
   persist(
     immer((set, get, ...rest) => ({
+      isRehydrated: false,
       api: null,
       setApi: (api) =>
         set((state) => {
           state.api = api;
         }),
-
       drawerOpen: false,
       openDrawer: () =>
         set((state) => {
@@ -61,6 +61,7 @@ const useBaseStore = create<
         }),
       selectedLogs: createSelectedLogsSlice(set, get, ...rest),
       logFilters: createLogFiltersSlice(set, get, ...rest),
+      columnVisibility: createColumnVisibilitySlice(set, get, ...rest),
     })),
     persistOptions,
   ),

app/src/utils/hooks.ts

@@ -205,3 +205,12 @@ export const useTagNames = () => {
     { enabled: !!selectedProjectId },
   );
 };
+
+export const useIsClientRehydrated = () => {
+  const isRehydrated = useAppStore((state) => state.isRehydrated);
+  const [isMounted, setIsMounted] = useState(false);
+  useEffect(() => {
+    setIsMounted(true);
+  }, []);
+  return isRehydrated && isMounted;
+};