Add useful datasets (#213)

* Create dataset from request logs

* Move drawer expansion logic out of app state

* Add empty dataset page

* Properly handle zero dataset state

* Add DatasetEntriesTable

* Open DatasetEntryEditorDrawer on row click

* Add editable messages

* Change Request Logs link to be a span

* Add FunctionCallEditor

* Change styling around

* Stop logging variant stats after a while

* Change FunctionCallEditor widths

* Record input tokens even on errored calls

* Allow user to add messages

* Allow changing from empty text to function call

* Fix some data layout issues

* Default to empty output

* Update arguments on blur

* Add beta flag to datasets tab

* Remove unused import

* Save training and testing datasets on fine tune

* Add DatasetEntryType

* Condense migrations

* Add index to datasetEntry

* Add datasetEntry index

* Fix types

* Enable scrolling beyond last line in VariantEditor

* Divide new dataset entries exactly along training/testing ratio

app/@types/nextjs-routes.d.ts

@@ -19,6 +19,8 @@ declare module "nextjs-routes" {
     | DynamicRoute<"/api/v1/[...trpc]", { "trpc": string[] }>
     | StaticRoute<"/api/v1/openapi">
     | StaticRoute<"/dashboard">
+    | DynamicRoute<"/datasets/[id]", { "id": string }>
+    | StaticRoute<"/datasets">
     | DynamicRoute<"/experiments/[experimentSlug]", { "experimentSlug": string }>
     | StaticRoute<"/experiments">
     | StaticRoute<"/fine-tunes">

app/prisma/migrations/20230904234505_revamp_dataset_entry/migration.sql

@@ -0,0 +1,27 @@
+/*
+  Warnings:
+
+  - Added the required column `input` to the `DatasetEntry` table without a default value. This is not possible if the table is not empty.
+  - Added the required column `inputTokens` to the `DatasetEntry` table without a default value. This is not possible if the table is not empty.
+  - Added the required column `outputTokens` to the `DatasetEntry` table without a default value. This is not possible if the table is not empty.
+  - Added the required column `type` to the `DatasetEntry` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- CreateEnum
+CREATE TYPE "DatasetEntryType" AS ENUM ('TRAIN', 'TEST');
+
+-- AlterTable
+ALTER TABLE "Dataset" ADD COLUMN     "trainingRatio" DOUBLE PRECISION NOT NULL DEFAULT 0.8;
+
+-- AlterTable
+ALTER TABLE "DatasetEntry" ADD COLUMN     "input" JSONB NOT NULL,
+ADD COLUMN     "inputTokens" INTEGER NOT NULL,
+ADD COLUMN     "output" JSONB,
+ADD COLUMN     "outputTokens" INTEGER NOT NULL,
+ADD COLUMN     "type" "DatasetEntryType" NOT NULL;
+
+-- CreateIndex
+CREATE INDEX "DatasetEntry_datasetId_createdAt_id_idx" ON "DatasetEntry"("datasetId", "createdAt", "id");
+
+-- CreateIndex
+CREATE INDEX "DatasetEntry_datasetId_type_idx" ON "DatasetEntry"("datasetId", "type");

app/prisma/schema.prisma

@@ -179,9 +179,10 @@ model OutputEvaluation {
 model Dataset {
     id String @id @default(uuid()) @db.Uuid
 
-    name           String
+    name                String
-    datasetEntries DatasetEntry[]
+    datasetEntries      DatasetEntry[]
-    fineTunes      FineTune[]
+    fineTunes           FineTune[]
+    trainingRatio  Float @default(0.8)
 
     projectId String  @db.Uuid
     project   Project @relation(fields: [projectId], references: [id], onDelete: Cascade)
@@ -190,17 +191,32 @@ model Dataset {
     updatedAt DateTime @updatedAt
 }
 
+enum DatasetEntryType {
+    TRAIN
+    TEST
+}
+
 model DatasetEntry {
     id String @id @default(uuid()) @db.Uuid
 
     loggedCallId String @db.Uuid
     loggedCall   LoggedCall @relation(fields: [loggedCallId], references: [id], onDelete: Cascade)
 
+    input        Json
+    output       Json?
+    inputTokens  Int
+    outputTokens Int
+
+    type DatasetEntryType
+
     datasetId String   @db.Uuid
     dataset   Dataset? @relation(fields: [datasetId], references: [id], onDelete: Cascade)
 
     createdAt DateTime @default(now())
     updatedAt DateTime @updatedAt
+
+    @@index([datasetId, createdAt, id])
+    @@index([datasetId, type])
 }
 
 model Project {
@@ -452,7 +468,7 @@ model FineTune {
     deploymentFinishedAt DateTime?
 
     datasetId String   @db.Uuid
-    dataset   Dataset @relation(fields: [datasetId], references: [id], onDelete: Cascade)
+    dataset    Dataset @relation(fields: [datasetId], references: [id], onDelete: Cascade)
 
     projectId String  @db.Uuid
     project   Project @relation(fields: [projectId], references: [id], onDelete: Cascade)

app/prisma/seedAgiEval.ts

@@ -1,5 +1,4 @@
 import { prisma } from "~/server/db";
-import { generateNewCell } from "~/server/utils/generateNewCell";
 import dedent from "dedent";
 import { execSync } from "child_process";
 import fs from "fs";

app/prisma/seedDashboard.ts

@@ -108,7 +108,7 @@ const MODEL_RESPONSE_TEMPLATES: {
     inputTokens: 236,
     outputTokens: 5,
     finishReason: "stop",
-    tags: [{ name: "prompt_id", value: "define_func" }],
+    tags: [{ name: "prompt_id", value: "add_scenario" }],
   },
   {
     reqPayload: {
@@ -311,7 +311,7 @@ const MODEL_RESPONSE_TEMPLATES: {
     outputTokens: 108,
     finishReason: "stop",
     tags: [
-      { name: "prompt_id", value: "chatcmpl-7" },
+      { name: "prompt_id", value: "define_func" },
       { name: "some_other_tag", value: "some_other_value" },
     ],
   },

app/src/components/ActionButton.tsx

@@ -3,7 +3,7 @@ import { useState } from "react";
 import { Button, HStack, type ButtonProps, Icon, Text } from "@chakra-ui/react";
 import { type IconType } from "react-icons";
 import { useAppStore } from "~/state/store";
-import { BetaModal } from "../BetaModal";
+import { BetaModal } from "./BetaModal";
 
 const ActionButton = ({
   icon,

app/src/components/InputDropdown.tsx

@@ -16,12 +16,16 @@ import {
 
 import { FiChevronDown } from "react-icons/fi";
 import { BiCheck } from "react-icons/bi";
+import { isEqual } from "lodash-es";
+import React from "react";
 
 type InputDropdownProps<T> = {
   options: ReadonlyArray<T>;
   selectedOption: T;
   onSelect: (option: T) => void;
   inputGroupProps?: InputGroupProps;
+  getDisplayLabel?: (option: T) => string;
+  isDisabled?: boolean;
 };
 
 const InputDropdown = <T,>({
@@ -29,19 +33,21 @@ const InputDropdown = <T,>({
   selectedOption,
   onSelect,
   inputGroupProps,
+  getDisplayLabel = (option) => option as string,
+  isDisabled,
 }: InputDropdownProps<T>) => {
-  const popover = useDisclosure();
+  const { onOpen, ...popover } = useDisclosure();
 
   return (
-    <Popover placement="bottom-start" {...popover}>
+    <Popover placement="bottom-start" onOpen={isDisabled ? undefined : onOpen} {...popover}>
       <PopoverTrigger>
         <InputGroup
           cursor="pointer"
-          w={(selectedOption as string).length * 14 + 180}
+          w={getDisplayLabel(selectedOption).length * 14 + 180}
           {...inputGroupProps}
         >
           <Input
-            value={selectedOption as string}
+            value={getDisplayLabel(selectedOption)}
             // eslint-disable-next-line @typescript-eslint/no-empty-function -- controlled input requires onChange
             onChange={() => {}}
             cursor="pointer"
@@ -52,9 +58,10 @@ const InputDropdown = <T,>({
             onFocus={(e) => {
               e.target.blur();
             }}
+            isDisabled={isDisabled}
           />
           <InputRightElement>
-            <Icon as={FiChevronDown} />
+            <Icon as={FiChevronDown} color={isDisabled ? "gray.300" : undefined} />
           </InputRightElement>
         </InputGroup>
       </PopoverTrigger>
@@ -78,8 +85,10 @@ const InputDropdown = <T,>({
               fontSize="sm"
               borderBottomWidth={1}
             >
-              <Text mr={16}>{option as string}</Text>
+              <Text mr={16}>{getDisplayLabel(option)}</Text>
-              {option === selectedOption && <Icon as={BiCheck} color="blue.500" boxSize={5} />}
+              {isEqual(option, selectedOption) && (
+                <Icon as={BiCheck} color="blue.500" boxSize={5} />
+              )}
             </HStack>
           ))}
         </VStack>

app/src/components/OutputsTable/ScenariosHeader.tsx

@@ -19,15 +19,13 @@ import {
   useScenarios,
 } from "~/utils/hooks";
 import { BsGear, BsPencil, BsPlus, BsStars } from "react-icons/bs";
-import { useAppStore } from "~/state/store";
 import { api } from "~/utils/api";
 
 export const ActionButton = (props: ButtonProps) => (
   <Button size="sm" variant="ghost" color="gray.600" {...props} />
 );
 
-export const ScenariosHeader = () => {
+export const ScenariosHeader = ({ openDrawer }: { openDrawer: () => void }) => {
-  const openDrawer = useAppStore((s) => s.openDrawer);
   const { canModify } = useExperimentAccess();
   const scenarios = useScenarios();
 

app/src/components/OutputsTable/VariantStats.tsx

@@ -20,7 +20,7 @@ export default function VariantStats(props: { variant: PromptVariant }) {
         inputTokens: 0,
         outputTokens: 0,
         scenarioCount: 0,
-        outputCount: 0,
+        finishedCount: 0,
         awaitingCompletions: false,
         awaitingEvals: false,
       },
@@ -42,7 +42,7 @@ export default function VariantStats(props: { variant: PromptVariant }) {
 
   const scale = chroma.scale([failColor, neutralColor, passColor]).domain([0, 0.5, 1]);
 
-  const showNumFinished = data.scenarioCount > 0 && data.scenarioCount !== data.outputCount;
+  const showNumFinished = data.scenarioCount > 0 && data.scenarioCount !== data.finishedCount;
 
   return (
     <HStack
@@ -55,7 +55,7 @@ export default function VariantStats(props: { variant: PromptVariant }) {
       <HStack px={cellPadding.x} flexWrap="wrap">
         {showNumFinished && (
           <Text>
-            {data.outputCount} / {data.scenarioCount}
+            {data.finishedCount} / {data.scenarioCount}
           </Text>
         )}
         {data.evalResults.map((result) => {

app/src/components/OutputsTable/index.tsx

@@ -12,7 +12,13 @@ import ScenarioPaginator from "./ScenarioPaginator";
 import { Fragment } from "react";
 import useScrolledPast from "./useHasScrolledPast";
 
-export default function OutputsTable({ experimentId }: { experimentId: string | undefined }) {
+export default function OutputsTable({
+  experimentId,
+  openDrawer,
+}: {
+  experimentId: string | undefined;
+  openDrawer: () => void;
+}) {
   const variants = api.promptVariants.list.useQuery(
     { experimentId: experimentId as string },
     { enabled: !!experimentId },
@@ -91,7 +97,7 @@ export default function OutputsTable({ experimentId }: { experimentId: string |
         colStart={1}
         borderRightWidth={0}
       >
-        <ScenariosHeader />
+        <ScenariosHeader openDrawer={openDrawer} />
       </GridItem>
 
       {scenarios.data.scenarios.map((scenario, i) => (

app/src/components/datasets/DatasetConfigurationDrawer/DatasetConfigurationDrawer.tsx

@@ -0,0 +1,37 @@
+import {
+  Drawer,
+  DrawerBody,
+  DrawerCloseButton,
+  DrawerContent,
+  DrawerHeader,
+  DrawerOverlay,
+  Heading,
+  VStack,
+  type UseDisclosureReturn,
+} from "@chakra-ui/react";
+
+import { DeleteButton } from "./DeleteButton";
+
+export default function DatasetConfigurationDrawer({
+  disclosure,
+}: {
+  disclosure: UseDisclosureReturn;
+}) {
+  return (
+    <Drawer placement="right" size="md" {...disclosure}>
+      <DrawerOverlay />
+      <DrawerContent>
+        <DrawerCloseButton />
+        <DrawerHeader>
+          <Heading size="md">Dataset Configuration</Heading>
+        </DrawerHeader>
+        <DrawerBody h="full" pb={4}>
+          <VStack h="full" justifyContent="space-between">
+            <VStack spacing={6}></VStack>
+            <DeleteButton closeDrawer={disclosure.onClose} />
+          </VStack>
+        </DrawerBody>
+      </DrawerContent>
+    </Drawer>
+  );
+}

app/src/components/datasets/DatasetConfigurationDrawer/DeleteButton.tsx

@@ -0,0 +1,39 @@
+import { Button, Icon, useDisclosure, Text } from "@chakra-ui/react";
+import { useRouter } from "next/router";
+import { BsTrash } from "react-icons/bs";
+
+import { useHandledAsyncCallback, useDataset } from "~/utils/hooks";
+import DeleteDatasetDialog from "./DeleteDatasetDialog";
+
+export const DeleteButton = ({ closeDrawer }: { closeDrawer: () => void }) => {
+  const dataset = useDataset();
+  const router = useRouter();
+
+  const disclosure = useDisclosure();
+
+  const [onDelete] = useHandledAsyncCallback(async () => {
+    await router.push({ pathname: "/datasets" });
+    closeDrawer();
+  }, [router, closeDrawer]);
+
+  return (
+    <>
+      <Button
+        size="sm"
+        variant="ghost"
+        colorScheme="red"
+        fontWeight="normal"
+        onClick={disclosure.onOpen}
+      >
+        <Icon as={BsTrash} boxSize={4} />
+        <Text ml={2}>Delete Dataset</Text>
+      </Button>
+
+      <DeleteDatasetDialog
+        datasetId={dataset.data?.id}
+        onDelete={onDelete}
+        disclosure={disclosure}
+      />
+    </>
+  );
+};

app/src/components/datasets/DatasetConfigurationDrawer/DeleteDatasetDialog.tsx

@@ -0,0 +1,73 @@
+import { useRef } from "react";
+import {
+  type UseDisclosureReturn,
+  AlertDialog,
+  AlertDialogOverlay,
+  AlertDialogContent,
+  AlertDialogHeader,
+  AlertDialogBody,
+  AlertDialogFooter,
+  Button,
+} from "@chakra-ui/react";
+import { api } from "~/utils/api";
+
+import { useHandledAsyncCallback } from "~/utils/hooks";
+
+const DeleteDatasetDialog = ({
+  datasetId,
+  onDelete,
+  disclosure,
+}: {
+  datasetId?: string;
+  onDelete?: () => void;
+  disclosure: UseDisclosureReturn;
+}) => {
+  const cancelRef = useRef<HTMLButtonElement>(null);
+
+  const mutation = api.datasets.delete.useMutation();
+  const utils = api.useContext();
+
+  const [onDeleteConfirm, deletionInProgress] = useHandledAsyncCallback(async () => {
+    if (!datasetId) return;
+    await mutation.mutateAsync({ id: datasetId });
+    await utils.datasets.list.invalidate();
+    onDelete?.();
+
+    disclosure.onClose();
+  }, [mutation, datasetId, disclosure.onClose]);
+
+  console.log("dataset id", datasetId);
+
+  return (
+    <AlertDialog leastDestructiveRef={cancelRef} {...disclosure}>
+      <AlertDialogOverlay>
+        <AlertDialogContent>
+          <AlertDialogHeader fontSize="lg" fontWeight="bold">
+            Delete Dataset
+          </AlertDialogHeader>
+
+          <AlertDialogBody>
+            If you delete this dataset all the associated dataset entries will be deleted as well.
+            Are you sure?
+          </AlertDialogBody>
+
+          <AlertDialogFooter>
+            <Button ref={cancelRef} onClick={disclosure.onClose}>
+              Cancel
+            </Button>
+            <Button
+              colorScheme="red"
+              isLoading={deletionInProgress}
+              onClick={onDeleteConfirm}
+              ml={3}
+            >
+              Delete
+            </Button>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialogOverlay>
+    </AlertDialog>
+  );
+};
+
+export default DeleteDatasetDialog;

app/src/components/datasets/DatasetEntriesTable/DatasetEntriesTable.tsx

@@ -0,0 +1,46 @@
+import { Card, Table, Tbody } from "@chakra-ui/react";
+import { useState } from "react";
+import { useDatasetEntries } from "~/utils/hooks";
+import { TableHeader, TableRow, EmptyTableRow } from "./TableRow";
+import DatasetEntryEditorDrawer from "./DatasetEntryEditorDrawer";
+
+export default function DatasetEntriesTable() {
+  const [expandedDatasetEntryId, setExpandedDatasetEntryId] = useState<string | null>(null);
+  const datasetEntries = useDatasetEntries().data?.entries;
+
+  return (
+    <>
+      <Card width="100%" overflowX="auto">
+        <Table>
+          <TableHeader />
+          <Tbody>
+            {datasetEntries?.length ? (
+              datasetEntries?.map((entry) => {
+                return (
+                  <TableRow
+                    key={entry.id}
+                    datasetEntry={entry}
+                    onToggle={() => {
+                      if (entry.id === expandedDatasetEntryId) {
+                        setExpandedDatasetEntryId(null);
+                      } else {
+                        setExpandedDatasetEntryId(entry.id);
+                      }
+                    }}
+                    showOptions
+                  />
+                );
+              })
+            ) : (
+              <EmptyTableRow />
+            )}
+          </Tbody>
+        </Table>
+      </Card>
+      <DatasetEntryEditorDrawer
+        datasetEntryId={expandedDatasetEntryId}
+        clearDatasetEntryId={() => setExpandedDatasetEntryId(null)}
+      />
+    </>
+  );
+}

app/src/components/datasets/DatasetEntriesTable/DatasetEntryEditorDrawer.tsx

@@ -0,0 +1,174 @@
+import { useState, useEffect, useMemo } from "react";
+import {
+  Drawer,
+  DrawerBody,
+  DrawerCloseButton,
+  DrawerContent,
+  DrawerHeader,
+  DrawerOverlay,
+  DrawerFooter,
+  Heading,
+  VStack,
+  HStack,
+  Button,
+  Text,
+  Divider,
+  Icon,
+} from "@chakra-ui/react";
+import { type CreateChatCompletionRequestMessage } from "openai/resources/chat";
+import { BsPlus } from "react-icons/bs";
+import { type DatasetEntryType } from "@prisma/client";
+
+import { api } from "~/utils/api";
+import { useDatasetEntry, useHandledAsyncCallback } from "~/utils/hooks";
+import EditableMessage from "./EditableMessage";
+import EntryTypeDropdown from "./EntryTypeDropdown";
+
+export default function DatasetDentryEditorDrawer({
+  datasetEntryId,
+  clearDatasetEntryId,
+}: {
+  datasetEntryId: string | null;
+  clearDatasetEntryId: () => void;
+}) {
+  const utils = api.useContext();
+
+  const datasetEntry = useDatasetEntry(datasetEntryId).data;
+
+  const savedInputMessages = useMemo(
+    () => datasetEntry?.input as unknown as CreateChatCompletionRequestMessage[],
+    [datasetEntry],
+  );
+  const savedOutputMessage = useMemo(
+    () => datasetEntry?.output as unknown as CreateChatCompletionRequestMessage,
+    [datasetEntry],
+  );
+
+  const [inputMessagesToSave, setInputMessagesToSave] = useState<
+    CreateChatCompletionRequestMessage[]
+  >([]);
+  const [outputMessageToSave, setOutputMessageToSave] =
+    useState<CreateChatCompletionRequestMessage | null>(null);
+
+  useEffect(() => {
+    if (savedInputMessages) {
+      setInputMessagesToSave(savedInputMessages);
+      setOutputMessageToSave(savedOutputMessage);
+    }
+  }, [savedInputMessages, savedOutputMessage]);
+
+  const updateMutation = api.datasetEntries.update.useMutation();
+  const [onSave, savingInProgress] = useHandledAsyncCallback(async () => {
+    if (!datasetEntryId || !inputMessagesToSave) return;
+    await updateMutation.mutateAsync({
+      id: datasetEntryId,
+      updates: {
+        input: JSON.stringify(inputMessagesToSave),
+        output: JSON.stringify(outputMessageToSave),
+      },
+    });
+    await utils.datasetEntries.list.invalidate();
+    await utils.datasetEntries.get.invalidate({ id: datasetEntryId });
+  }, [updateMutation, datasetEntryId, inputMessagesToSave, outputMessageToSave, utils]);
+
+  const [onUpdateType] = useHandledAsyncCallback(
+    async (type: DatasetEntryType) => {
+      if (!datasetEntryId) return;
+      await updateMutation.mutateAsync({
+        id: datasetEntryId,
+        updates: {
+          type,
+        },
+      });
+      await utils.datasetEntries.list.invalidate();
+      await utils.datasetEntries.get.invalidate({ id: datasetEntryId });
+    },
+    [updateMutation, datasetEntryId, utils],
+  );
+
+  return (
+    <Drawer isOpen={!!datasetEntryId} onClose={clearDatasetEntryId} placement="right" size="md">
+      <DrawerOverlay />
+      <DrawerContent>
+        <DrawerCloseButton pt={6} />
+        <DrawerHeader bgColor="orange.50">
+          <HStack w="full" justifyContent="space-between" pr={8}>
+            <Heading size="md">Dataset Entry</Heading>
+            {datasetEntry && (
+              <EntryTypeDropdown type={datasetEntry.type} onTypeChange={onUpdateType} />
+            )}
+          </HStack>
+        </DrawerHeader>
+        <DrawerBody h="full" pb={4} bgColor="orange.50">
+          <VStack h="full" justifyContent="space-between">
+            <VStack w="full" spacing={12} py={4}>
+              <VStack w="full" alignItems="flex-start">
+                <Text fontWeight="bold">Input</Text>
+                {inputMessagesToSave.map((message, i) => {
+                  return (
+                    <>
+                      <Divider key={`divider-${i}`} my={4} />
+                      <EditableMessage
+                        key={i}
+                        message={message}
+                        onEdit={(message) => {
+                          const newInputMessages = [...inputMessagesToSave];
+                          newInputMessages[i] = message;
+                          setInputMessagesToSave(newInputMessages);
+                        }}
+                        onDelete={() => {
+                          const newInputMessages = [...inputMessagesToSave];
+                          newInputMessages.splice(i, 1);
+                          setInputMessagesToSave(newInputMessages);
+                        }}
+                      />
+                    </>
+                  );
+                })}
+                <Divider my={4} />
+                <Button
+                  w="full"
+                  onClick={() =>
+                    setInputMessagesToSave([...inputMessagesToSave, { role: "user", content: "" }])
+                  }
+                  variant="outline"
+                  color="gray.500"
+                  _hover={{ bgColor: "orange.100" }}
+                >
+                  <HStack spacing={0}>
+                    <Text>Add Message</Text>
+                    <Icon as={BsPlus} boxSize={6} />
+                  </HStack>
+                </Button>
+              </VStack>
+              <VStack w="full" alignItems="flex-start">
+                <Text fontWeight="bold">Output</Text>
+                <Divider my={4} />
+                <EditableMessage
+                  message={outputMessageToSave}
+                  onEdit={(message) => setOutputMessageToSave(message)}
+                  isOutput
+                />
+              </VStack>
+            </VStack>
+          </VStack>
+        </DrawerBody>
+        <DrawerFooter bgColor="orange.50">
+          <HStack>
+            <Button
+              onClick={() => {
+                setInputMessagesToSave(savedInputMessages);
+                setOutputMessageToSave(savedOutputMessage);
+              }}
+            >
+              Reset
+            </Button>
+            <Button isLoading={savingInProgress} onClick={onSave} colorScheme="orange">
+              Save
+            </Button>
+          </HStack>
+        </DrawerFooter>
+      </DrawerContent>
+    </Drawer>
+  );
+}

app/src/components/datasets/DatasetEntriesTable/EditableMessage.tsx

@@ -0,0 +1,105 @@
+import { VStack, HStack, Tooltip, IconButton, Icon } from "@chakra-ui/react";
+import { type CreateChatCompletionRequestMessage } from "openai/resources/chat";
+import { BsX } from "react-icons/bs";
+
+import AutoResizeTextArea from "~/components/AutoResizeTextArea";
+import InputDropdown from "~/components/InputDropdown";
+import { parseableToFunctionCall } from "~/utils/utils";
+import FunctionCallEditor from "./FunctionCallEditor";
+
+const MESSAGE_ROLE_OPTIONS = ["system", "user", "assistant", "function"] as const;
+const OUTPUT_OPTIONS = ["plaintext", "func_call"] as const;
+
+const EditableMessage = ({
+  message,
+  onEdit,
+  onDelete,
+  isOutput,
+}: {
+  message: CreateChatCompletionRequestMessage | null;
+  onEdit: (message: CreateChatCompletionRequestMessage) => void;
+  onDelete?: () => void;
+  isOutput?: boolean;
+}) => {
+  const { role = "assistant", content = "", function_call } = message || {};
+
+  const currentOutputOption: (typeof OUTPUT_OPTIONS)[number] = function_call
+    ? "func_call"
+    : "plaintext";
+
+  return (
+    <VStack w="full">
+      <HStack w="full" justifyContent="space-between">
+        <HStack>
+          {!isOutput && (
+            <InputDropdown
+              options={MESSAGE_ROLE_OPTIONS}
+              selectedOption={role}
+              onSelect={(option) => {
+                const updatedMessage = { role: option, content };
+                if (role === "assistant" && currentOutputOption === "func_call") {
+                  updatedMessage.content = JSON.stringify(function_call, null, 2);
+                }
+                onEdit(updatedMessage);
+              }}
+              inputGroupProps={{ w: "32", bgColor: "white" }}
+            />
+          )}
+          {role === "assistant" && (
+            <InputDropdown
+              options={OUTPUT_OPTIONS}
+              selectedOption={currentOutputOption}
+              onSelect={(option) => {
+                const updatedMessage: CreateChatCompletionRequestMessage = {
+                  role,
+                  content: null,
+                  function_call: undefined,
+                };
+                if (option === "plaintext") {
+                  updatedMessage.content = JSON.stringify(function_call, null, 2);
+                } else if (option === "func_call") {
+                  updatedMessage.function_call =
+                    content && parseableToFunctionCall(content)
+                      ? JSON.parse(content)
+                      : { name: "", arguments: "{}" };
+                }
+                onEdit(updatedMessage);
+              }}
+              inputGroupProps={{ w: "32", bgColor: "white" }}
+            />
+          )}
+        </HStack>
+        {!isOutput && (
+          <HStack>
+            <Tooltip label="Delete" hasArrow>
+              <IconButton
+                aria-label="Delete"
+                icon={<Icon as={BsX} boxSize={6} />}
+                onClick={onDelete}
+                size="xs"
+                display="flex"
+                colorScheme="gray"
+                color="gray.500"
+                variant="ghost"
+              />
+            </Tooltip>
+          </HStack>
+        )}
+      </HStack>
+      {function_call ? (
+        <FunctionCallEditor
+          function_call={function_call}
+          onEdit={(function_call) => onEdit({ role, function_call, content: null })}
+        />
+      ) : (
+        <AutoResizeTextArea
+          value={content || JSON.stringify(function_call, null, 2)}
+          onChange={(e) => onEdit({ role, content: e.target.value })}
+          bgColor="white"
+        />
+      )}
+    </VStack>
+  );
+};
+
+export default EditableMessage;

app/src/components/datasets/DatasetEntriesTable/EntryTypeDropdown.tsx

@@ -0,0 +1,24 @@
+import { type DatasetEntryType } from "@prisma/client";
+
+import InputDropdown from "~/components/InputDropdown";
+
+const ENTRY_TYPE_OPTIONS: DatasetEntryType[] = ["TRAIN", "TEST"];
+
+const EntryTypeDropdown = ({
+  type,
+  onTypeChange,
+}: {
+  type: DatasetEntryType;
+  onTypeChange: (type: DatasetEntryType) => void;
+}) => {
+  return (
+    <InputDropdown
+      options={ENTRY_TYPE_OPTIONS}
+      selectedOption={type}
+      onSelect={onTypeChange}
+      inputGroupProps={{ w: "32", bgColor: "white" }}
+    />
+  );
+};
+
+export default EntryTypeDropdown;

app/src/components/datasets/DatasetEntriesTable/FunctionCallEditor.tsx

@@ -0,0 +1,125 @@
+import { useRef, useMemo, useEffect } from "react";
+import { VStack, HStack, Text, Input, Box } from "@chakra-ui/react";
+import { type CreateChatCompletionRequestMessage } from "openai/resources/chat";
+
+import { useAppStore } from "~/state/store";
+import { type CreatedEditor } from "~/state/sharedVariantEditor.slice";
+
+const FunctionCallEditor = ({
+  function_call,
+  onEdit,
+}: {
+  function_call: CreateChatCompletionRequestMessage.FunctionCall;
+  onEdit: (function_call: CreateChatCompletionRequestMessage.FunctionCall) => void;
+}) => {
+  const monaco = useAppStore.use.sharedArgumentsEditor.monaco();
+  const editorRef = useRef<CreatedEditor | null>(null);
+  const editorId = useMemo(() => `editor_${Math.random().toString(36).substring(7)}`, []);
+
+  useEffect(() => {
+    if (monaco) {
+      const container = document.getElementById(editorId) as HTMLElement;
+
+      const editor = monaco.editor.create(container, {
+        value: function_call.arguments,
+        language: "json",
+        theme: "customTheme",
+        lineNumbers: "off",
+        minimap: { enabled: false },
+        wrappingIndent: "indent",
+        wrappingStrategy: "advanced",
+        wordWrap: "on",
+        folding: false,
+        scrollbar: {
+          alwaysConsumeMouseWheel: false,
+          verticalScrollbarSize: 0,
+        },
+        wordWrapBreakAfterCharacters: "",
+        wordWrapBreakBeforeCharacters: "",
+        quickSuggestions: true,
+        renderLineHighlight: "none",
+        fontSize: 14,
+        scrollBeyondLastLine: false,
+      });
+
+      editorRef.current = editor;
+
+      const updateHeight = () => {
+        const contentHeight = editor.getContentHeight();
+        container.style.height = `${contentHeight}px`;
+        editor.layout();
+      };
+
+      const attemptDocumentFormat = () => {
+        const action = editor.getAction("editor.action.formatDocument");
+        if (action) {
+          action
+            .run()
+            .then(updateHeight)
+            .catch((error) => {
+              console.error("Error running formatDocument:", error);
+            });
+          return true;
+        }
+        return false;
+      };
+
+      editor.onDidBlurEditorText(() => {
+        attemptDocumentFormat();
+        onEdit({ name: function_call.name, arguments: editor.getValue() });
+      });
+
+      // Interval function to check for action availability
+      const checkForActionInterval = setInterval(() => {
+        const formatted = attemptDocumentFormat();
+        if (formatted) {
+          clearInterval(checkForActionInterval); // Clear the interval once the action is found and run
+        }
+      }, 100); // Check every 100ms
+
+      // Add content change listener
+      const contentChangeListener = editor.onDidChangeModelContent(updateHeight);
+
+      const resizeObserver = new ResizeObserver(() => {
+        editor.layout();
+      });
+      resizeObserver.observe(container);
+
+      return () => {
+        contentChangeListener.dispose();
+        resizeObserver.disconnect();
+        editor?.dispose();
+      };
+    }
+  }, [monaco, editorId, function_call.name, function_call.arguments, onEdit]);
+
+  return (
+    <VStack w="full" alignItems="flex-start">
+      <HStack w="full">
+        <Text fontWeight="bold" w={192}>
+          Name:
+        </Text>
+        <Input
+          value={function_call.name}
+          onChange={(e) => onEdit({ name: e.target.value, arguments: function_call.arguments })}
+          bgColor="white"
+        />
+      </HStack>
+      <Text fontWeight="bold" w={32}>
+        Arguments
+      </Text>
+      <VStack
+        borderRadius={4}
+        border="1px solid"
+        borderColor="gray.200"
+        w="full"
+        py={1}
+        bgColor="white"
+      >
+        <Box id={editorId} w="full" />
+      </VStack>
+    </VStack>
+  );
+};
+
+export default FunctionCallEditor;

app/src/components/datasets/DatasetEntriesTable/TableRow.tsx

@@ -0,0 +1,128 @@
+import { Box, Td, Tr, Thead, Th, Tooltip, HStack, Text, Checkbox } from "@chakra-ui/react";
+import Link from "next/link";
+
+import dayjs from "~/utils/dayjs";
+import { type RouterOutputs } from "~/utils/api";
+import { useAppStore } from "~/state/store";
+import { useIsClientRehydrated, useDatasetEntries } from "~/utils/hooks";
+import { useMemo } from "react";
+
+type DatasetEntry = RouterOutputs["datasetEntries"]["list"]["entries"][0];
+
+export const TableHeader = () => {
+  const matchingDatasetEntryIds = useDatasetEntries().data?.matchingEntryIds;
+  const selectedDatasetEntryIds = useAppStore((s) => s.selectedDatasetEntries.selectedIds);
+  const addSelectedIds = useAppStore((s) => s.selectedDatasetEntries.addSelectedIds);
+  const clearSelectedIds = useAppStore((s) => s.selectedDatasetEntries.clearSelectedIds);
+  const allSelected = useMemo(() => {
+    if (!matchingDatasetEntryIds || !matchingDatasetEntryIds.length) return false;
+    return matchingDatasetEntryIds.every((id) => selectedDatasetEntryIds.has(id));
+  }, [matchingDatasetEntryIds, selectedDatasetEntryIds]);
+  const isClientRehydrated = useIsClientRehydrated();
+  if (!isClientRehydrated) return null;
+
+  return (
+    <Thead>
+      <Tr>
+        <Th pr={0}>
+          <HStack minW={16}>
+            <Checkbox
+              isChecked={allSelected}
+              onChange={() => {
+                allSelected ? clearSelectedIds() : addSelectedIds(matchingDatasetEntryIds || []);
+              }}
+            />
+            <Text>
+              ({selectedDatasetEntryIds.size ? `${selectedDatasetEntryIds.size}/` : ""}
+              {matchingDatasetEntryIds?.length || 0})
+            </Text>
+          </HStack>
+        </Th>
+        <Th>Created At</Th>
+        <Th isNumeric>Input tokens</Th>
+        <Th isNumeric>Output tokens</Th>
+        <Th isNumeric>Type</Th>
+      </Tr>
+    </Thead>
+  );
+};
+
+export const TableRow = ({
+  datasetEntry,
+  onToggle,
+  showOptions,
+}: {
+  datasetEntry: DatasetEntry;
+  onToggle: () => void;
+  showOptions?: boolean;
+}) => {
+  const createdAt = dayjs(datasetEntry.createdAt).format("MMMM D h:mm A");
+  const fullTime = dayjs(datasetEntry.createdAt).toString();
+
+  const isChecked = useAppStore((s) => s.selectedDatasetEntries.selectedIds.has(datasetEntry.id));
+  const toggleChecked = useAppStore((s) => s.selectedDatasetEntries.toggleSelectedId);
+
+  const isClientRehydrated = useIsClientRehydrated();
+  if (!isClientRehydrated) return null;
+
+  return (
+    <Tr
+      onClick={onToggle}
+      key={datasetEntry.id}
+      _hover={{ bgColor: "gray.50", cursor: "pointer" }}
+      fontSize="sm"
+    >
+      {showOptions && (
+        <Td>
+          <Checkbox isChecked={isChecked} onChange={() => toggleChecked(datasetEntry.id)} />
+        </Td>
+      )}
+      <Td>
+        <Tooltip label={fullTime} placement="top">
+          <Box whiteSpace="nowrap" minW="120px">
+            {createdAt}
+          </Box>
+        </Tooltip>
+      </Td>
+      <Td isNumeric>{datasetEntry.inputTokens}</Td>
+      <Td isNumeric>{datasetEntry.outputTokens}</Td>
+      <Td isNumeric>{datasetEntry.type}</Td>
+    </Tr>
+  );
+};
+
+export const EmptyTableRow = ({ filtersApplied = true }: { filtersApplied?: boolean }) => {
+  const visibleColumns = useAppStore((s) => s.columnVisibility.visibleColumns);
+  const filters = useAppStore((state) => state.logFilters.filters);
+  const { isLoading } = useDatasetEntries();
+
+  if (isLoading) return null;
+
+  if (filters.length && filtersApplied) {
+    return (
+      <Tr>
+        <Td w="full" colSpan={visibleColumns.size + 1}>
+          <Text color="gray.500" textAlign="center" w="full" p={4}>
+            No matching entries found. Try removing some filters.
+          </Text>
+        </Td>
+      </Tr>
+    );
+  }
+
+  return (
+    <Tr>
+      <Td w="full" colSpan={visibleColumns.size + 1}>
+        <Text color="gray.500" textAlign="center" w="full" p={4}>
+          This dataset has no entries. Add some logs in the{" "}
+          <Link href="/request-logs">
+            <Text as="span" color="blue.600">
+              Request Logs
+            </Text>
+          </Link>{" "}
+          tab.
+        </Text>
+      </Td>
+    </Tr>
+  );
+};

app/src/components/datasets/DatasetEntryPaginator.tsx

@@ -0,0 +1,16 @@
+import { type StackProps } from "@chakra-ui/react";
+
+import { useDatasetEntries } from "~/utils/hooks";
+import Paginator from "../Paginator";
+
+const DatasetEntryPaginator = (props: StackProps) => {
+  const { data } = useDatasetEntries();
+
+  if (!data) return null;
+
+  const { matchingEntryIds } = data;
+
+  return <Paginator count={matchingEntryIds.length} {...props} />;
+};
+
+export default DatasetEntryPaginator;

app/src/components/datasets/DatasetHeaderButtons.tsx

@@ -0,0 +1,20 @@
+import { Button, HStack, Icon, Text } from "@chakra-ui/react";
+import { useDataset } from "~/utils/hooks";
+import { BsGearFill } from "react-icons/bs";
+
+export const DatasetHeaderButtons = ({ openDrawer }: { openDrawer: () => void }) => {
+  const dataset = useDataset();
+
+  if (dataset.isLoading) return null;
+
+  return (
+    <HStack spacing={0} mt={{ base: 2, md: 0 }}>
+      <Button variant={{ base: "solid", md: "ghost" }} onClick={openDrawer}>
+        <HStack>
+          <Icon as={BsGearFill} />
+          <Text>Configure</Text>
+        </HStack>
+      </Button>
+    </HStack>
+  );
+};

app/src/components/datasets/DatasetsTable.tsx

@@ -0,0 +1,52 @@
+import { Card, Table, Thead, Tr, Th, Tbody, Td, VStack, Icon, Text } from "@chakra-ui/react";
+import { FaTable } from "react-icons/fa";
+import Link from "next/link";
+
+import dayjs from "~/utils/dayjs";
+import { useDatasets } from "~/utils/hooks";
+
+const DatasetsTable = ({}) => {
+  const { data } = useDatasets();
+
+  const datasets = data || [];
+
+  return (
+    <Card width="100%" overflowX="auto">
+      {datasets.length ? (
+        <Table>
+          <Thead>
+            <Tr>
+              <Th>Name</Th>
+              <Th>Created At</Th>
+              <Th>Size</Th>
+            </Tr>
+          </Thead>
+          <Tbody>
+            {datasets.map((dataset) => {
+              return (
+                <Tr key={dataset.id}>
+                  <Td>
+                    <Link href={{ pathname: "/datasets/[id]", query: { id: dataset.id } }}>
+                      <Text color="blue.600">{dataset.name}</Text>
+                    </Link>
+                  </Td>
+                  <Td>{dayjs(dataset.createdAt).format("MMMM D h:mm A")}</Td>
+                  <Td>{dataset._count.datasetEntries}</Td>
+                </Tr>
+              );
+            })}
+          </Tbody>
+        </Table>
+      ) : (
+        <VStack py={8}>
+          <Icon as={FaTable} boxSize={16} color="gray.300" />
+          <Text color="gray.400" fontSize="lg" fontWeight="bold">
+            No Datasets Found. Create your first dataset.
+          </Text>
+        </VStack>
+      )}
+    </Card>
+  );
+};
+
+export default DatasetsTable;

app/src/components/datasets/ExperimentButton.tsx

@@ -0,0 +1,21 @@
+import { RiFlaskLine } from "react-icons/ri";
+
+import { useAppStore } from "~/state/store";
+import ActionButton from "../ActionButton";
+
+const ExperimentButton = () => {
+  const selectedIds = useAppStore((s) => s.selectedDatasetEntries.selectedIds);
+  return (
+    <ActionButton
+      onClick={() => {
+        console.log("experimenting with these ids", selectedIds);
+      }}
+      label="Experiment"
+      icon={RiFlaskLine}
+      isDisabled={selectedIds.size === 0}
+      requireBeta
+    />
+  );
+};
+
+export default ExperimentButton;

app/src/components/datasets/FineTuneButton.tsx

@@ -20,17 +20,19 @@ import { AiTwotoneThunderbolt } from "react-icons/ai";
 import humanId from "human-id";
 import { useRouter } from "next/router";
 
-import { useHandledAsyncCallback } from "~/utils/hooks";
+import { useDataset, useDatasetEntries, useHandledAsyncCallback } from "~/utils/hooks";
 import { api } from "~/utils/api";
 import { useAppStore } from "~/state/store";
-import ActionButton from "./ActionButton";
+import ActionButton from "../ActionButton";
 import InputDropdown from "../InputDropdown";
 import { FiChevronDown } from "react-icons/fi";
 
 const SUPPORTED_BASE_MODELS = ["llama2-7b", "llama2-13b", "llama2-70b", "gpt-3.5-turbo"];
 
 const FineTuneButton = () => {
-  const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
+  const datasetEntries = useDatasetEntries().data;
+
+  const numEntries = datasetEntries?.matchingEntryIds.length || 0;
 
   const disclosure = useDisclosure();
 
@@ -40,7 +42,7 @@ const FineTuneButton = () => {
         onClick={disclosure.onOpen}
         label="Fine Tune"
         icon={AiTwotoneThunderbolt}
-        isDisabled={selectedLogIds.size === 0}
+        isDisabled={numEntries === 0}
         requireBeta
       />
       <FineTuneModal disclosure={disclosure} />
@@ -52,8 +54,8 @@ export default FineTuneButton;
 
 const FineTuneModal = ({ disclosure }: { disclosure: UseDisclosureReturn }) => {
   const selectedProjectId = useAppStore((s) => s.selectedProjectId);
-  const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
+  const dataset = useDataset().data;
-  const clearSelectedLogIds = useAppStore((s) => s.selectedLogs.clearSelectedLogIds);
+  const datasetEntries = useDatasetEntries().data;
 
   const [selectedBaseModel, setSelectedBaseModel] = useState(SUPPORTED_BASE_MODELS[0]);
   const [modelSlug, setModelSlug] = useState(humanId({ separator: "-", capitalize: false }));
@@ -71,19 +73,17 @@ const FineTuneModal = ({ disclosure }: { disclosure: UseDisclosureReturn }) => {
   const createFineTuneMutation = api.fineTunes.create.useMutation();
 
   const [createFineTune, creationInProgress] = useHandledAsyncCallback(async () => {
-    if (!selectedProjectId || !modelSlug || !selectedBaseModel || !selectedLogIds.size) return;
+    if (!selectedProjectId || !modelSlug || !selectedBaseModel || !dataset) return;
     await createFineTuneMutation.mutateAsync({
-      projectId: selectedProjectId,
       slug: modelSlug,
       baseModel: selectedBaseModel,
-      selectedLogIds: Array.from(selectedLogIds),
+      datasetId: dataset.id,
     });
 
     await utils.fineTunes.list.invalidate();
     await router.push({ pathname: "/fine-tunes" });
-    clearSelectedLogIds();
     disclosure.onClose();
-  }, [createFineTuneMutation, selectedProjectId, selectedLogIds, modelSlug, selectedBaseModel]);
+  }, [createFineTuneMutation, selectedProjectId, modelSlug, selectedBaseModel]);
 
   return (
     <Modal size={{ base: "xl", md: "2xl" }} {...disclosure}>
@@ -99,7 +99,8 @@ const FineTuneModal = ({ disclosure }: { disclosure: UseDisclosureReturn }) => {
         <ModalBody maxW="unset">
           <VStack w="full" spacing={8} pt={4} alignItems="flex-start">
             <Text>
-              We'll train on the <b>{selectedLogIds.size}</b> logs you've selected.
+              We'll train on <b>{datasetEntries?.trainingCount}</b> and test on{" "}
+              <b>{datasetEntries?.testingCount}</b> entries in this dataset.
             </Text>
             <VStack>
               <HStack spacing={2} w="full">

app/src/components/experiments/DeleteExperimentDialog.tsx

@@ -27,7 +27,7 @@ const DeleteExperimentDialog = ({
   const mutation = api.experiments.delete.useMutation();
   const utils = api.useContext();
 
-  const [onDeleteConfirm] = useHandledAsyncCallback(async () => {
+  const [onDeleteConfirm, deletionInProgress] = useHandledAsyncCallback(async () => {
     if (!experimentId) return;
     await mutation.mutateAsync({ id: experimentId });
     await utils.experiments.list.invalidate();
@@ -53,7 +53,12 @@ const DeleteExperimentDialog = ({
             <Button ref={cancelRef} onClick={disclosure.onClose}>
               Cancel
             </Button>
-            <Button colorScheme="red" onClick={onDeleteConfirm} ml={3}>
+            <Button
+              colorScheme="red"
+              isLoading={deletionInProgress}
+              onClick={onDeleteConfirm}
+              ml={3}
+            >
               Delete
             </Button>
           </AlertDialogFooter>

app/src/components/experiments/ExperimentHeaderButtons/DeleteDialog.tsx

@@ -1,57 +0,0 @@
-import {
-  Button,
-  AlertDialog,
-  AlertDialogBody,
-  AlertDialogFooter,
-  AlertDialogHeader,
-  AlertDialogContent,
-  AlertDialogOverlay,
-} from "@chakra-ui/react";
-
-import { useRouter } from "next/router";
-import { useRef } from "react";
-import { api } from "~/utils/api";
-import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
-
-export const DeleteDialog = ({ onClose }: { onClose: () => void }) => {
-  const experiment = useExperiment();
-  const deleteMutation = api.experiments.delete.useMutation();
-  const utils = api.useContext();
-  const router = useRouter();
-
-  const cancelRef = useRef<HTMLButtonElement>(null);
-
-  const [onDeleteConfirm] = useHandledAsyncCallback(async () => {
-    if (!experiment.data?.id) return;
-    await deleteMutation.mutateAsync({ id: experiment.data.id });
-    await utils.experiments.list.invalidate();
-    await router.push({ pathname: "/experiments" });
-    onClose();
-  }, [deleteMutation, experiment.data?.id, router]);
-
-  return (
-    <AlertDialog isOpen leastDestructiveRef={cancelRef} onClose={onClose}>
-      <AlertDialogOverlay>
-        <AlertDialogContent>
-          <AlertDialogHeader fontSize="lg" fontWeight="bold">
-            Delete Experiment
-          </AlertDialogHeader>
-
-          <AlertDialogBody>
-            If you delete this experiment all the associated prompts and scenarios will be deleted
-            as well. Are you sure?
-          </AlertDialogBody>
-
-          <AlertDialogFooter>
-            <Button ref={cancelRef} onClick={onClose}>
-              Cancel
-            </Button>
-            <Button colorScheme="red" onClick={onDeleteConfirm} ml={3}>
-              Delete
-            </Button>
-          </AlertDialogFooter>
-        </AlertDialogContent>
-      </AlertDialogOverlay>
-    </AlertDialog>
-  );
-};

app/src/components/experiments/ExperimentHeaderButtons/ExperimentHeaderButtons.tsx

@@ -3,17 +3,14 @@ import { useOnForkButtonPressed } from "./useOnForkButtonPressed";
 import { useExperiment } from "~/utils/hooks";
 import { BsGearFill } from "react-icons/bs";
 import { TbGitFork } from "react-icons/tb";
-import { useAppStore } from "~/state/store";
 
-export const ExperimentHeaderButtons = () => {
+export const ExperimentHeaderButtons = ({ openDrawer }: { openDrawer: () => void }) => {
   const experiment = useExperiment();
 
   const canModify = experiment.data?.access.canModify ?? false;
 
   const { onForkButtonPressed, isForking } = useOnForkButtonPressed();
 
-  const openDrawer = useAppStore((s) => s.openDrawer);
-
   if (experiment.isLoading) return null;
 
   return (

app/src/components/experiments/ExperimentSettingsDrawer/DeleteButton.tsx

@@ -2,17 +2,15 @@ import { Button, Icon, useDisclosure, Text } from "@chakra-ui/react";
 import { useRouter } from "next/router";
 import { BsTrash } from "react-icons/bs";
 
-import { useAppStore } from "~/state/store";
 import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
-import DeleteExperimentDialog from "../experiments/DeleteExperimentDialog";
+import DeleteExperimentDialog from "../DeleteExperimentDialog";
 
-export const DeleteButton = () => {
+export const DeleteButton = ({ closeDrawer }: { closeDrawer: () => void }) => {
   const experiment = useExperiment();
   const router = useRouter();
 
   const disclosure = useDisclosure();
 
-  const closeDrawer = useAppStore((s) => s.closeDrawer);
   const [onDelete] = useHandledAsyncCallback(async () => {
     await router.push({ pathname: "/experiments" });
     closeDrawer();

app/src/components/experiments/ExperimentSettingsDrawer/EditEvaluations.tsx

@@ -19,7 +19,7 @@ import { useCallback, useState } from "react";
 import { BsPencil, BsX } from "react-icons/bs";
 import { api } from "~/utils/api";
 import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
-import AutoResizeTextArea from "../AutoResizeTextArea";
+import AutoResizeTextArea from "~/components/AutoResizeTextArea";
 
 type EvalValues = Pick<Evaluation, "label" | "value" | "evalType">;
 

app/src/components/experiments/ExperimentSettingsDrawer/EditScenarioVars.tsx

@@ -5,7 +5,7 @@ import { BsPencil, BsX } from "react-icons/bs";
 import { api } from "~/utils/api";
 import { useExperiment, useHandledAsyncCallback, useScenarioVars } from "~/utils/hooks";
 import { maybeReportError } from "~/utils/errorHandling/maybeReportError";
-import { FloatingLabelInput } from "./FloatingLabelInput";
+import { FloatingLabelInput } from "~/components/OutputsTable/FloatingLabelInput";
 
 export const ScenarioVar = ({
   variable,

app/src/components/experiments/ExperimentSettingsDrawer/ExperimentSettingsDrawer.tsx

@@ -7,18 +7,19 @@ import {
   DrawerOverlay,
   Heading,
   VStack,
+  type UseDisclosureReturn,
 } from "@chakra-ui/react";
-import EditScenarioVars from "../OutputsTable/EditScenarioVars";
+import EditScenarioVars from "./EditScenarioVars";
-import EditEvaluations from "../OutputsTable/EditEvaluations";
+import EditEvaluations from "./EditEvaluations";
-import { useAppStore } from "~/state/store";
 import { DeleteButton } from "./DeleteButton";
 
-export default function ExperimentSettingsDrawer() {
+export default function ExperimentSettingsDrawer({
-  const isOpen = useAppStore((state) => state.drawerOpen);
+  disclosure,
-  const closeDrawer = useAppStore((state) => state.closeDrawer);
+}: {
-
+  disclosure: UseDisclosureReturn;
+}) {
   return (
-    <Drawer isOpen={isOpen} placement="right" onClose={closeDrawer} size="md">
+    <Drawer placement="right" size="md" {...disclosure}>
       <DrawerOverlay />
       <DrawerContent>
         <DrawerCloseButton />
@@ -31,7 +32,7 @@ export default function ExperimentSettingsDrawer() {
               <EditScenarioVars />
               <EditEvaluations />
             </VStack>
-            <DeleteButton />
+            <DeleteButton closeDrawer={disclosure.onClose} />
           </VStack>
         </DrawerBody>
       </DrawerContent>

app/src/components/nav/AppShell.tsx

@@ -17,7 +17,7 @@ import { useRouter } from "next/router";
 import { BsGearFill, BsGithub, BsPersonCircle } from "react-icons/bs";
 import { IoStatsChartOutline } from "react-icons/io5";
 import { RiHome3Line, RiFlaskLine } from "react-icons/ri";
-import { AiOutlineThunderbolt } from "react-icons/ai";
+import { AiOutlineThunderbolt, AiOutlineDatabase } from "react-icons/ai";
 import { FaReadme } from "react-icons/fa";
 import { signIn, useSession } from "next-auth/react";
 
@@ -78,6 +78,7 @@ const NavSidebar = () => {
 
             <IconLink icon={RiHome3Line} label="Dashboard" href="/dashboard" />
             <IconLink icon={IoStatsChartOutline} label="Request Logs" href="/request-logs" />
+            <IconLink icon={AiOutlineDatabase} label="Datasets" href="/datasets" beta />
             <IconLink icon={AiOutlineThunderbolt} label="Fine Tunes" href="/fine-tunes" beta />
             <IconLink icon={RiFlaskLine} label="Experiments" href="/experiments" />
             <VStack w="full" alignItems="flex-start" spacing={0} pt={8}>

app/src/components/requestLogs/AddToDatasetButton.tsx

@@ -0,0 +1,194 @@
+import { useState, useEffect, useMemo } from "react";
+import {
+  Modal,
+  ModalOverlay,
+  ModalContent,
+  ModalHeader,
+  ModalCloseButton,
+  ModalBody,
+  ModalFooter,
+  HStack,
+  VStack,
+  Icon,
+  Text,
+  Button,
+  Flex,
+  Input,
+  useDisclosure,
+  type UseDisclosureReturn,
+  Checkbox,
+} from "@chakra-ui/react";
+import { FiPlusSquare } from "react-icons/fi";
+
+import { useDatasets, useHandledAsyncCallback } from "~/utils/hooks";
+import { api } from "~/utils/api";
+import { useAppStore } from "~/state/store";
+import ActionButton from "../ActionButton";
+import InputDropdown from "../InputDropdown";
+import { maybeReportError } from "~/utils/errorHandling/maybeReportError";
+import { useRouter } from "next/router";
+
+const AddToDatasetButton = () => {
+  const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
+
+  const disclosure = useDisclosure();
+
+  return (
+    <>
+      <ActionButton
+        onClick={disclosure.onOpen}
+        label="Add to Dataset"
+        icon={FiPlusSquare}
+        isDisabled={selectedLogIds.size === 0}
+        requireBeta
+      />
+      <AddToDatasetModal disclosure={disclosure} />
+    </>
+  );
+};
+
+export default AddToDatasetButton;
+
+const AddToDatasetModal = ({ disclosure }: { disclosure: UseDisclosureReturn }) => {
+  const selectedProjectId = useAppStore((s) => s.selectedProjectId);
+  const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
+  const clearSelectedLogIds = useAppStore((s) => s.selectedLogs.clearSelectedLogIds);
+  const router = useRouter();
+
+  const datasets = useDatasets().data;
+
+  const existingDatasetOptions = useMemo(
+    () =>
+      datasets?.length
+        ? datasets.map((d) => ({ label: d.name, id: d.id }))
+        : [{ label: "", id: "" }],
+    [datasets],
+  );
+
+  const [selectedDatasetOption, setSelectedDatasetOption] = useState(existingDatasetOptions?.[0]);
+  const [newDatasetName, setNewDatasetName] = useState("");
+  const [createNewDataset, setCreateNewDataset] = useState(false);
+
+  useEffect(() => {
+    if (disclosure.isOpen) {
+      setSelectedDatasetOption(existingDatasetOptions?.[0]);
+      setCreateNewDataset(!existingDatasetOptions[0]?.id);
+    }
+  }, [disclosure.isOpen, existingDatasetOptions]);
+
+  const createDatasetEntriesMutation = api.datasetEntries.create.useMutation();
+
+  const [addToDataset, addingInProgress] = useHandledAsyncCallback(async () => {
+    if (
+      !selectedProjectId ||
+      !selectedLogIds.size ||
+      !(createNewDataset ? newDatasetName : selectedDatasetOption?.id)
+    )
+      return;
+    const datasetParams = createNewDataset
+      ? { newDatasetParams: { projectId: selectedProjectId, name: newDatasetName } }
+      : { datasetId: selectedDatasetOption?.id };
+    const response = await createDatasetEntriesMutation.mutateAsync({
+      loggedCallIds: Array.from(selectedLogIds),
+      ...datasetParams,
+    });
+
+    if (maybeReportError(response)) return;
+
+    const datasetId = response.payload;
+
+    await router.push({ pathname: "/datasets/[id]", query: { id: datasetId } });
+
+    disclosure.onClose();
+    clearSelectedLogIds();
+  }, [
+    selectedProjectId,
+    selectedLogIds,
+    createNewDataset,
+    selectedDatasetOption?.id,
+    newDatasetName,
+    router,
+  ]);
+
+  return (
+    <Modal size={{ base: "xl", md: "2xl" }} {...disclosure}>
+      <ModalOverlay />
+      <ModalContent w={1200}>
+        <ModalHeader>
+          <HStack>
+            <Icon as={FiPlusSquare} />
+            <Text>Add to Dataset</Text>
+          </HStack>
+        </ModalHeader>
+        <ModalCloseButton />
+        <ModalBody maxW="unset">
+          <VStack w="full" spacing={8} pt={4} alignItems="flex-start">
+            <Text>
+              We'll add the <b>{selectedLogIds.size}</b> logs you have selected to the dataset you
+              choose.
+            </Text>
+            <VStack alignItems="flex-start" spacing={4}>
+              {existingDatasetOptions?.length && selectedDatasetOption && (
+                <Flex
+                  flexDir={{ base: "column", md: "row" }}
+                  alignItems={{ base: "flex-start", md: "center" }}
+                >
+                  <Text fontWeight="bold" w={48}>
+                    Dataset:
+                  </Text>
+                  <InputDropdown
+                    options={existingDatasetOptions}
+                    selectedOption={selectedDatasetOption}
+                    getDisplayLabel={(option) => option.label}
+                    onSelect={(option) => setSelectedDatasetOption(option)}
+                    inputGroupProps={{ w: 48 }}
+                    isDisabled={createNewDataset}
+                  />
+                  <Checkbox
+                    isChecked={createNewDataset}
+                    onChange={(e) => setCreateNewDataset(e.target.checked)}
+                    paddingLeft={4}
+                    isDisabled={!existingDatasetOptions[0]?.id}
+                  >
+                    <Text>Create New Dataset</Text>
+                  </Checkbox>
+                </Flex>
+              )}
+
+              {createNewDataset && (
+                <Flex
+                  flexDir={{ base: "column", md: "row" }}
+                  alignItems={{ base: "flex-start", md: "center" }}
+                >
+                  <Text w={48} fontWeight="bold">
+                    Dataset Name:
+                  </Text>
+                  <Input
+                    w={48}
+                    value={newDatasetName}
+                    onChange={(e) => setNewDatasetName(e.target.value)}
+                  />
+                </Flex>
+              )}
+            </VStack>
+          </VStack>
+        </ModalBody>
+        <ModalFooter>
+          <HStack>
+            <Button colorScheme="gray" onClick={disclosure.onClose} minW={24}>
+              Cancel
+            </Button>
+            <Button
+              colorScheme="blue"
+              onClick={addToDataset}
+              isLoading={addingInProgress}
+              minW={24}
+            >
+              Add
+            </Button>
+          </HStack>
+        </ModalFooter>
+      </ModalContent>
+    </Modal>
+  );
+};

app/src/components/requestLogs/ColumnVisiblityDropdown.tsx

@@ -17,7 +17,7 @@ import { useMemo } from "react";
 import { useIsClientRehydrated, useTagNames } from "~/utils/hooks";
 import { useAppStore } from "~/state/store";
 import { StaticColumnKeys } from "~/state/columnVisiblitySlice";
-import ActionButton from "./ActionButton";
+import ActionButton from "../ActionButton";
 
 const ColumnVisiblityDropdown = () => {
   const tagNames = useTagNames().data;

app/src/components/requestLogs/ExportButton.tsx

@@ -28,7 +28,7 @@ import { BiExport } from "react-icons/bi";
 import { useHandledAsyncCallback } from "~/utils/hooks";
 import { api } from "~/utils/api";
 import { useAppStore } from "~/state/store";
-import ActionButton from "./ActionButton";
+import ActionButton from "../ActionButton";
 import InputDropdown from "../InputDropdown";
 import { FiChevronUp, FiChevronDown } from "react-icons/fi";
 import InfoCircle from "../InfoCircle";
@@ -81,7 +81,7 @@ const ExportLogsModal = ({ disclosure }: { disclosure: UseDisclosureReturn }) =>
       return;
     const response = await exportLogsMutation.mutateAsync({
       projectId: selectedProjectId,
-      selectedLogIds: Array.from(selectedLogIds),
+      loggedCallIds: Array.from(selectedLogIds),
       testingSplit,
       selectedExportFormat,
       removeDuplicates,

app/src/components/requestLogs/TableRow.tsx

@@ -9,17 +9,14 @@ import {
   Collapse,
   HStack,
   VStack,
-  Button,
-  ButtonGroup,
   Text,
   Checkbox,
   Link as ChakraLink,
 } from "@chakra-ui/react";
-import Link from "next/link";
 
 import dayjs from "~/utils/dayjs";
 import { type RouterOutputs } from "~/utils/api";
-import { FormattedJson } from "./FormattedJson";
+import { FormattedJson } from "../FormattedJson";
 import { useAppStore } from "~/state/store";
 import { useIsClientRehydrated, useLoggedCalls, useTagNames } from "~/utils/hooks";
 import { useMemo } from "react";
@@ -176,23 +173,16 @@ export const TableRow = ({
       <Tr>
         <Td colSpan={visibleColumns.size + 1} w="full" p={0}>
           <Collapse in={isExpanded} unmountOnExit={true}>
-            <VStack p={4} align="stretch">
+            <HStack align="stretch" p={4}>
-              <HStack align="stretch">
+              <VStack flex={1} align="stretch">
-                <VStack flex={1} align="stretch">
+                <Heading size="sm">Input</Heading>
-                  <Heading size="sm">Input</Heading>
+                <FormattedJson json={loggedCall.modelResponse?.reqPayload} />
-                  <FormattedJson json={loggedCall.modelResponse?.reqPayload} />
+              </VStack>
-                </VStack>
+              <VStack flex={1} align="stretch">
-                <VStack flex={1} align="stretch">
+                <Heading size="sm">Output</Heading>
-                  <Heading size="sm">Output</Heading>
+                <FormattedJson json={loggedCall.modelResponse?.respPayload} />
-                  <FormattedJson json={loggedCall.modelResponse?.respPayload} />
+              </VStack>
-                </VStack>
+            </HStack>
-              </HStack>
-              <ButtonGroup alignSelf="flex-end">
-                <Button as={Link} colorScheme="blue" href={{ pathname: "/experiments" }}>
-                  Experiments
-                </Button>
-              </ButtonGroup>
-            </VStack>
           </Collapse>
         </Td>
       </Tr>

app/src/modelProviders/openai-ChatCompletion/index.ts

@@ -42,24 +42,21 @@ const modelProvider: OpenaiChatModelProvider = {
   canStream: true,
   getCompletion,
   getUsage: (input, output) => {
-    if (output.choices.length === 0) return null;
-
     const model = modelProvider.getModel(input);
     if (!model) return null;
 
     let inputTokens: number;
     let outputTokens: number;
 
-    if (output.usage) {
+    if (output?.usage) {
       inputTokens = output.usage.prompt_tokens;
       outputTokens = output.usage.completion_tokens;
     } else {
       try {
         inputTokens = countOpenAIChatTokens(model, input.messages);
-        outputTokens = countOpenAIChatTokens(
+        outputTokens = output
-          model,
+          ? countOpenAIChatTokens(model, output.choices.map((c) => c.message).filter(truthyFilter))
-          output.choices.map((c) => c.message).filter(truthyFilter),
+          : 0;
-        );
       } catch (err) {
         inputTokens = 0;
         outputTokens = 0;

app/src/modelProviders/types.ts

@@ -59,7 +59,7 @@ export type ModelProvider<SupportedModels extends string, InputSchema, OutputSch
   ) => Promise<CompletionResponse<OutputSchema>>;
   getUsage: (
     input: InputSchema,
-    output: OutputSchema,
+    output?: OutputSchema,
   ) => { gpuRuntime?: number; inputTokens?: number; outputTokens?: number; cost?: number } | null;
 
   // This is just a convenience for type inference, don't use it at runtime

app/src/pages/datasets/[id].tsx

@@ -0,0 +1,113 @@
+import {
+  Breadcrumb,
+  BreadcrumbItem,
+  Center,
+  Flex,
+  Icon,
+  Input,
+  VStack,
+  HStack,
+  useDisclosure,
+} from "@chakra-ui/react";
+import Link from "next/link";
+
+import { useState, useEffect } from "react";
+import { AiOutlineDatabase } from "react-icons/ai";
+import AppShell from "~/components/nav/AppShell";
+import { api } from "~/utils/api";
+import { useDataset, useHandledAsyncCallback } from "~/utils/hooks";
+import PageHeaderContainer from "~/components/nav/PageHeaderContainer";
+import ProjectBreadcrumbContents from "~/components/nav/ProjectBreadcrumbContents";
+import DatasetConfigurationDrawer from "~/components/datasets/DatasetConfigurationDrawer/DatasetConfigurationDrawer";
+import { DatasetHeaderButtons } from "~/components/datasets/DatasetHeaderButtons";
+import DatasetEntriesTable from "~/components/datasets/DatasetEntriesTable/DatasetEntriesTable";
+import DatasetEntryPaginator from "~/components/datasets/DatasetEntryPaginator";
+import { useAppStore } from "~/state/store";
+import FineTuneButton from "~/components/datasets/FineTuneButton";
+import ExperimentButton from "~/components/datasets/ExperimentButton";
+
+export default function Dataset() {
+  const utils = api.useContext();
+
+  const dataset = useDataset();
+
+  const drawerDisclosure = useDisclosure();
+  const [name, setName] = useState(dataset.data?.name || "");
+  useEffect(() => {
+    setName(dataset.data?.name || "");
+  }, [dataset.data?.name]);
+
+  useEffect(() => {
+    useAppStore.getState().sharedArgumentsEditor.loadMonaco().catch(console.error);
+  }, []);
+
+  const updateMutation = api.datasets.update.useMutation();
+  const [onSaveName] = useHandledAsyncCallback(async () => {
+    if (name && name !== dataset.data?.name && dataset.data?.id) {
+      await updateMutation.mutateAsync({
+        id: dataset.data.id,
+        name,
+      });
+      await Promise.all([utils.datasets.list.invalidate(), utils.datasets.get.invalidate()]);
+    }
+  }, [updateMutation, dataset.data?.id, dataset.data?.name, name]);
+
+  if (!dataset.isLoading && !dataset.data) {
+    return (
+      <AppShell title="Dataset not found">
+        <Center h="100%">
+          <div>Dataset not found 😕</div>
+        </Center>
+      </AppShell>
+    );
+  }
+
+  return (
+    <>
+      <AppShell title={dataset.data?.name}>
+        <VStack h="full" overflowY="scroll">
+          <PageHeaderContainer>
+            <Breadcrumb>
+              <BreadcrumbItem>
+                <ProjectBreadcrumbContents projectName={dataset.data?.project?.name} />
+              </BreadcrumbItem>
+              <BreadcrumbItem>
+                <Link href="/datasets">
+                  <Flex alignItems="center" _hover={{ textDecoration: "underline" }}>
+                    <Icon as={AiOutlineDatabase} boxSize={4} mr={2} /> Datasets
+                  </Flex>
+                </Link>
+              </BreadcrumbItem>
+              <BreadcrumbItem isCurrentPage>
+                <Input
+                  size="sm"
+                  value={name}
+                  onChange={(e) => setName(e.target.value)}
+                  onBlur={onSaveName}
+                  borderWidth={1}
+                  borderColor="transparent"
+                  fontSize={16}
+                  px={0}
+                  minW={{ base: 100, lg: 300 }}
+                  flex={1}
+                  _hover={{ borderColor: "gray.300" }}
+                  _focus={{ borderColor: "blue.500", outline: "none" }}
+                />
+              </BreadcrumbItem>
+            </Breadcrumb>
+            <DatasetHeaderButtons openDrawer={drawerDisclosure.onOpen} />
+          </PageHeaderContainer>
+          <VStack px={8} py={8} alignItems="flex-start" spacing={4} w="full">
+            <HStack w="full" justifyContent="flex-end">
+              <FineTuneButton />
+              <ExperimentButton />
+            </HStack>
+            <DatasetEntriesTable />
+            <DatasetEntryPaginator />
+          </VStack>
+        </VStack>
+      </AppShell>
+      <DatasetConfigurationDrawer disclosure={drawerDisclosure} />
+    </>
+  );
+}

app/src/pages/datasets/index.tsx

@@ -0,0 +1,17 @@
+import { VStack, Text, Divider } from "@chakra-ui/react";
+import AppShell from "~/components/nav/AppShell";
+import DatasetsTable from "~/components/datasets/DatasetsTable";
+
+export default function DatasetsPage() {
+  return (
+    <AppShell title="Datasets" requireAuth>
+      <VStack w="full" py={8} px={8} spacing={4} alignItems="flex-start">
+        <Text fontSize="2xl" fontWeight="bold">
+          Datasets
+        </Text>
+        <Divider />
+        <DatasetsTable />
+      </VStack>
+    </AppShell>
+  );
+}

app/src/pages/experiments/[experimentSlug].tsx

@@ -8,26 +8,25 @@ import {
   Input,
   Text,
   VStack,
+  useDisclosure,
 } from "@chakra-ui/react";
 import Link from "next/link";
 
-import { useRouter } from "next/router";
 import { useState, useEffect } from "react";
 import { RiFlaskLine } from "react-icons/ri";
 import OutputsTable from "~/components/OutputsTable";
-import ExperimentSettingsDrawer from "~/components/ExperimentSettingsDrawer/ExperimentSettingsDrawer";
+import ExperimentSettingsDrawer from "~/components/experiments/ExperimentSettingsDrawer/ExperimentSettingsDrawer";
+import { ExperimentHeaderButtons } from "~/components/experiments/ExperimentHeaderButtons/ExperimentHeaderButtons";
 import AppShell from "~/components/nav/AppShell";
 import { api } from "~/utils/api";
 import { useExperiment, useHandledAsyncCallback } from "~/utils/hooks";
 import { useAppStore } from "~/state/store";
 import { useSyncVariantEditor } from "~/state/sync";
-import { ExperimentHeaderButtons } from "~/components/experiments/ExperimentHeaderButtons/ExperimentHeaderButtons";
 import Head from "next/head";
 import PageHeaderContainer from "~/components/nav/PageHeaderContainer";
 import ProjectBreadcrumbContents from "~/components/nav/ProjectBreadcrumbContents";
 
 export default function Experiment() {
-  const router = useRouter();
   const utils = api.useContext();
   useSyncVariantEditor();
 
@@ -44,6 +43,7 @@ export default function Experiment() {
     useAppStore.getState().sharedVariantEditor.loadMonaco().catch(console.error);
   }, []);
 
+  const drawerDisclosure = useDisclosure();
   const [label, setLabel] = useState(experiment.data?.label || "");
   useEffect(() => {
     setLabel(experiment.data?.label || "");
@@ -121,11 +121,11 @@ export default function Experiment() {
                 )}
               </BreadcrumbItem>
             </Breadcrumb>
-            <ExperimentHeaderButtons />
+            <ExperimentHeaderButtons openDrawer={drawerDisclosure.onOpen} />
           </PageHeaderContainer>
-          <ExperimentSettingsDrawer />
+          <ExperimentSettingsDrawer disclosure={drawerDisclosure} />
           <Box w="100%" overflowX="auto" flex={1} id="output-container">
-            <OutputsTable experimentId={experiment.data?.id} />
+            <OutputsTable experimentId={experiment.data?.id} openDrawer={drawerDisclosure.onOpen} />
           </Box>
         </VStack>
       </AppShell>

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

@@ -4,14 +4,13 @@ import { Text, VStack, Divider, HStack, Box } from "@chakra-ui/react";
 import AppShell from "~/components/nav/AppShell";
 import LoggedCallTable from "~/components/requestLogs/LoggedCallsTable";
 import LoggedCallsPaginator from "~/components/requestLogs/LoggedCallsPaginator";
-import ActionButton from "~/components/requestLogs/ActionButton";
+import ActionButton from "~/components/ActionButton";
 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";
-import FineTuneButton from "~/components/requestLogs/FineTuneButton";
 import ExportButton from "~/components/requestLogs/ExportButton";
+import AddToDatasetButton from "~/components/requestLogs/AddToDatasetButton";
 
 export default function LoggedCalls() {
   const selectedLogIds = useAppStore((s) => s.selectedLogs.selectedLogIds);
@@ -27,16 +26,7 @@ export default function LoggedCalls() {
           </Text>
           <Divider />
           <HStack w="full" justifyContent="flex-end">
-            <FineTuneButton />
+            <AddToDatasetButton />
-            <ActionButton
-              onClick={() => {
-                console.log("experimenting with these ids", selectedLogIds);
-              }}
-              label="Experiment"
-              icon={RiFlaskLine}
-              isDisabled={selectedLogIds.size === 0}
-              requireBeta
-            />
             <ExportButton />
             <ColumnVisiblityDropdown />
             <ActionButton

app/src/server/api/external/v1Api.router.ts

@@ -119,10 +119,10 @@ export const v1ApiRouter = createOpenApiRouter({
 
       let usage;
       let model;
-      if (reqPayload.success && respPayload.success) {
+      if (reqPayload.success) {
         usage = modelProvider.getUsage(
           input.reqPayload as CompletionCreateParams,
-          input.respPayload as ChatCompletion,
+          respPayload.success ? (input.respPayload as ChatCompletion) : undefined,
         );
         model = reqPayload.data.model;
       }

app/src/server/api/root.router.ts

@@ -9,6 +9,8 @@ import { worldChampsRouter } from "./routers/worldChamps.router";
 import { projectsRouter } from "./routers/projects.router";
 import { dashboardRouter } from "./routers/dashboard.router";
 import { loggedCallsRouter } from "./routers/loggedCalls.router";
+import { datasetsRouter } from "./routers/datasets.router";
+import { datasetEntriesRouter } from "./routers/datasetEntries.router";
 import { fineTunesRouter } from "./routers/fineTunes.router";
 import { usersRouter } from "./routers/users.router";
 import { adminJobsRouter } from "./routers/adminJobs.router";
@@ -29,6 +31,8 @@ export const appRouter = createTRPCRouter({
   projects: projectsRouter,
   dashboard: dashboardRouter,
   loggedCalls: loggedCallsRouter,
+  datasets: datasetsRouter,
+  datasetEntries: datasetEntriesRouter,
   fineTunes: fineTunesRouter,
   users: usersRouter,
   adminJobs: adminJobsRouter,

app/src/server/api/routers/datasetEntries.router.ts

@@ -0,0 +1,296 @@
+import { type Prisma } from "@prisma/client";
+import { z } from "zod";
+import { v4 as uuidv4 } from "uuid";
+import {
+  type ChatCompletion,
+  type CompletionCreateParams,
+  type CreateChatCompletionRequestMessage,
+} from "openai/resources/chat";
+import { TRPCError } from "@trpc/server";
+import { shuffle } from "lodash-es";
+
+import { createTRPCRouter, protectedProcedure } from "~/server/api/trpc";
+import { prisma } from "~/server/db";
+import { requireCanModifyProject, requireCanViewProject } from "~/utils/accessControl";
+import { error, success } from "~/utils/errorHandling/standardResponses";
+import { countOpenAIChatTokens } from "~/utils/countTokens";
+
+export const datasetEntriesRouter = createTRPCRouter({
+  list: protectedProcedure
+    .input(z.object({ datasetId: z.string(), page: z.number(), pageSize: z.number() }))
+    .query(async ({ input, ctx }) => {
+      const { datasetId, page, pageSize } = input;
+
+      const { projectId } = await prisma.dataset.findUniqueOrThrow({
+        where: { id: datasetId },
+      });
+      await requireCanViewProject(projectId, ctx);
+
+      const [entries, matchingEntries, trainingCount, testingCount] = await prisma.$transaction([
+        prisma.datasetEntry.findMany({
+          where: {
+            datasetId: datasetId,
+          },
+          orderBy: [{ createdAt: "desc" }, { id: "desc" }],
+          skip: (page - 1) * pageSize,
+          take: pageSize,
+        }),
+        prisma.datasetEntry.findMany({
+          where: {
+            datasetId: datasetId,
+          },
+          select: {
+            id: true,
+          },
+        }),
+        prisma.datasetEntry.count({
+          where: {
+            datasetId: datasetId,
+            type: "TRAIN",
+          },
+        }),
+        prisma.datasetEntry.count({
+          where: {
+            datasetId: datasetId,
+            type: "TEST",
+          },
+        }),
+      ]);
+
+      return {
+        entries,
+        matchingEntryIds: matchingEntries.map((entry) => entry.id),
+        trainingCount,
+        testingCount,
+      };
+    }),
+  get: protectedProcedure.input(z.object({ id: z.string() })).query(async ({ input, ctx }) => {
+    const entry = await prisma.datasetEntry.findUniqueOrThrow({
+      where: { id: input.id },
+      include: {
+        dataset: true,
+      },
+    });
+
+    if (!entry.dataset) {
+      throw new TRPCError({ message: "Dataset not found for dataset entry", code: "NOT_FOUND" });
+    }
+
+    await requireCanViewProject(entry.dataset.projectId, ctx);
+
+    if (!entry) {
+      throw new TRPCError({ message: "Dataset entry not found", code: "NOT_FOUND" });
+    }
+
+    return entry;
+  }),
+  create: protectedProcedure
+    .input(
+      z.object({
+        datasetId: z.string().optional(),
+        newDatasetParams: z
+          .object({
+            projectId: z.string(),
+            name: z.string(),
+          })
+          .optional(),
+        loggedCallIds: z.string().array(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      let datasetId: string;
+      let trainingRatio = 0.8;
+      if (input.datasetId) {
+        datasetId = input.datasetId;
+        const { projectId, trainingRatio: datasetTrainingRatio } =
+          await prisma.dataset.findUniqueOrThrow({
+            where: { id: input.datasetId },
+          });
+        trainingRatio = datasetTrainingRatio;
+        await requireCanModifyProject(projectId, ctx);
+      } else if (input.newDatasetParams) {
+        await requireCanModifyProject(input.newDatasetParams.projectId, ctx);
+        datasetId = uuidv4();
+      } else {
+        return error("No datasetId or newDatasetParams provided");
+      }
+
+      const [loggedCalls, existingTrainingCount, existingTestingCount] = await prisma.$transaction([
+        prisma.loggedCall.findMany({
+          where: {
+            id: {
+              in: input.loggedCallIds,
+            },
+            modelResponse: {
+              isNot: null,
+            },
+          },
+          include: {
+            modelResponse: {
+              select: {
+                reqPayload: true,
+                respPayload: true,
+                inputTokens: true,
+                outputTokens: true,
+              },
+            },
+          },
+        }),
+        prisma.datasetEntry.count({
+          where: {
+            datasetId,
+            type: "TRAIN",
+          },
+        }),
+        prisma.datasetEntry.count({
+          where: {
+            datasetId,
+            type: "TEST",
+          },
+        }),
+      ]);
+
+      const shuffledLoggedCalls = shuffle(loggedCalls);
+
+      const totalEntries = existingTrainingCount + existingTestingCount + loggedCalls.length;
+      const numTrainingToAdd = Math.floor(trainingRatio * totalEntries) - existingTrainingCount;
+
+      const datasetEntriesToCreate: Prisma.DatasetEntryCreateManyInput[] = [];
+
+      let i = 0;
+      for (const loggedCall of shuffledLoggedCalls) {
+        const inputMessages = (
+          loggedCall.modelResponse?.reqPayload as unknown as CompletionCreateParams
+        ).messages;
+        let output: ChatCompletion.Choice.Message | undefined = undefined;
+        const resp = loggedCall.modelResponse?.respPayload as unknown as ChatCompletion | undefined;
+        if (resp && resp.choices?.[0]) {
+          output = resp.choices[0].message;
+        } else {
+          output = {
+            role: "assistant",
+            content: "",
+          };
+        }
+
+        datasetEntriesToCreate.push({
+          datasetId,
+          loggedCallId: loggedCall.id,
+          input: inputMessages as unknown as Prisma.InputJsonValue,
+          output: output as unknown as Prisma.InputJsonValue,
+          inputTokens: loggedCall.modelResponse?.inputTokens || 0,
+          outputTokens: loggedCall.modelResponse?.outputTokens || 0,
+          type: i < numTrainingToAdd ? "TRAIN" : "TEST",
+        });
+        i++;
+      }
+
+      // Ensure dataset and dataset entries are created atomically
+      await prisma.$transaction([
+        prisma.dataset.upsert({
+          where: { id: datasetId },
+          update: {},
+          create: {
+            id: datasetId,
+            projectId: input.newDatasetParams?.projectId ?? "",
+            name: input.newDatasetParams?.name ?? "",
+            trainingRatio,
+          },
+        }),
+        prisma.datasetEntry.createMany({
+          data: shuffle(datasetEntriesToCreate),
+        }),
+      ]);
+
+      return success(datasetId);
+    }),
+
+  update: protectedProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        updates: z.object({
+          type: z.enum(["TRAIN", "TEST"]).optional(),
+          input: z.string().optional(),
+          output: z.string().optional(),
+        }),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const { dataset } = await prisma.datasetEntry.findUniqueOrThrow({
+        where: { id: input.id },
+        include: {
+          dataset: true,
+        },
+      });
+
+      if (!dataset) {
+        return error("Dataset not found for dataset entry");
+      }
+
+      await requireCanModifyProject(dataset.projectId, ctx);
+
+      let parsedInput = undefined;
+      let inputTokens = undefined;
+      if (input.updates.input) {
+        parsedInput = JSON.parse(input.updates.input);
+        inputTokens = countOpenAIChatTokens(
+          "gpt-4-0613",
+          parsedInput as unknown as CreateChatCompletionRequestMessage[],
+        );
+      }
+
+      let parsedOutput = undefined;
+      let outputTokens = undefined;
+      if (input.updates.output) {
+        parsedOutput = JSON.parse(input.updates.output);
+        outputTokens = countOpenAIChatTokens("gpt-4-0613", [
+          parsedOutput as unknown as ChatCompletion.Choice.Message,
+        ]);
+      }
+
+      await prisma.datasetEntry.update({
+        where: { id: input.id },
+        data: {
+          type: input.updates.type,
+          input: parsedInput,
+          output: parsedOutput,
+          inputTokens,
+          outputTokens,
+        },
+      });
+
+      return success("Dataset entry updated");
+    }),
+
+  delete: protectedProcedure
+    .input(z.object({ ids: z.string().array() }))
+    .mutation(async ({ input, ctx }) => {
+      if (input.ids.length === 0) {
+        return error("No ids provided");
+      }
+      const { dataset } = await prisma.datasetEntry.findUniqueOrThrow({
+        where: { id: input.ids[0] },
+        include: {
+          dataset: true,
+        },
+      });
+
+      if (!dataset) {
+        return error("Dataset not found for dataset entry");
+      }
+
+      await requireCanModifyProject(dataset.projectId, ctx);
+
+      await prisma.datasetEntry.deleteMany({
+        where: {
+          id: {
+            in: input.ids,
+          },
+          datasetId: dataset?.id,
+        },
+      });
+
+      return success("Dataset entries deleted");
+    }),
+});

app/src/server/api/routers/datasets.router.ts

@@ -0,0 +1,97 @@
+import { z } from "zod";
+import { createTRPCRouter, protectedProcedure } from "~/server/api/trpc";
+import { prisma } from "~/server/db";
+import { requireCanModifyProject, requireCanViewProject } from "~/utils/accessControl";
+import { success } from "~/utils/errorHandling/standardResponses";
+
+export const datasetsRouter = createTRPCRouter({
+  get: protectedProcedure.input(z.object({ id: z.string() })).query(async ({ input, ctx }) => {
+    const dataset = await prisma.dataset.findUniqueOrThrow({
+      where: { id: input.id },
+      include: {
+        project: true,
+      },
+    });
+
+    await requireCanViewProject(dataset.projectId, ctx);
+
+    return dataset;
+  }),
+  list: protectedProcedure
+    .input(z.object({ projectId: z.string() }))
+    .query(async ({ input, ctx }) => {
+      await requireCanViewProject(input.projectId, ctx);
+
+      return await prisma.dataset.findMany({
+        where: {
+          projectId: input.projectId,
+        },
+        include: {
+          _count: {
+            select: {
+              datasetEntries: true,
+            },
+          },
+        },
+        orderBy: { createdAt: "desc" },
+      });
+    }),
+
+  create: protectedProcedure
+    .input(
+      z.object({
+        projectId: z.string(),
+        name: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      await requireCanModifyProject(input.projectId, ctx);
+
+      const dataset = await prisma.dataset.create({
+        data: {
+          projectId: input.projectId,
+          name: input.name,
+        },
+      });
+
+      return success(dataset);
+    }),
+
+  update: protectedProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        name: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const { projectId } = await prisma.dataset.findUniqueOrThrow({
+        where: { id: input.id },
+      });
+      await requireCanModifyProject(projectId, ctx);
+
+      await prisma.dataset.update({
+        where: { id: input.id },
+        data: {
+          name: input.name,
+        },
+      });
+
+      return success("Dataset updated");
+    }),
+
+  delete: protectedProcedure
+    .input(z.object({ id: z.string() }))
+    .mutation(async ({ input, ctx }) => {
+      const { projectId } = await prisma.dataset.findUniqueOrThrow({
+        where: { id: input.id },
+      });
+      await requireCanModifyProject(projectId, ctx);
+
+      await prisma.dataset.delete({
+        where: { id: input.id },
+      });
+
+      return success("Dataset deleted");
+    }),
+});

app/src/server/api/routers/fineTunes.router.ts

@@ -1,6 +1,4 @@
 import { z } from "zod";
-import { v4 as uuidv4 } from "uuid";
-import { type Prisma } from "@prisma/client";
 
 import { createTRPCRouter, protectedProcedure } from "~/server/api/trpc";
 import { prisma } from "~/server/db";
@@ -55,14 +53,18 @@ export const fineTunesRouter = createTRPCRouter({
   create: protectedProcedure
     .input(
       z.object({
-        projectId: z.string(),
+        datasetId: z.string(),
-        selectedLogIds: z.array(z.string()),
         slug: z.string(),
         baseModel: z.string(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      await requireCanModifyProject(input.projectId, ctx);
+      const { projectId } = await prisma.dataset.findUniqueOrThrow({
+        where: {
+          id: input.datasetId,
+        },
+      });
+      await requireCanModifyProject(projectId, ctx);
 
       const existingFineTune = await prisma.fineTune.findFirst({
         where: {
@@ -74,39 +76,14 @@ export const fineTunesRouter = createTRPCRouter({
         return error("A fine tune with that slug already exists");
       }
 
-      const newDatasetId = uuidv4();
+      await prisma.fineTune.create({
-
+        data: {
-      const datasetEntriesToCreate: Prisma.DatasetEntryCreateManyDatasetInput[] =
+          projectId,
-        input.selectedLogIds.map((loggedCallId) => ({
+          slug: input.slug,
-          loggedCallId,
+          baseModel: input.baseModel,
-        }));
+          datasetId: input.datasetId,
-
+        },
-      await prisma.$transaction([
+      });
-        prisma.dataset.create({
-          data: {
-            id: newDatasetId,
-            name: input.slug,
-            project: {
-              connect: {
-                id: input.projectId,
-              },
-            },
-            datasetEntries: {
-              createMany: {
-                data: datasetEntriesToCreate,
-              },
-            },
-          },
-        }),
-        prisma.fineTune.create({
-          data: {
-            projectId: input.projectId,
-            slug: input.slug,
-            baseModel: input.baseModel,
-            datasetId: newDatasetId,
-          },
-        }),
-      ]);
 
       return success();
     }),

app/src/server/api/routers/loggedCalls.router.ts

@@ -189,7 +189,7 @@ export const loggedCallsRouter = createTRPCRouter({
     .input(
       z.object({
         projectId: z.string(),
-        selectedLogIds: z.string().array(),
+        loggedCallIds: z.string().array(),
         testingSplit: z.number(),
         selectedExportFormat: z.string(),
         removeDuplicates: z.boolean(),
@@ -203,7 +203,7 @@ export const loggedCallsRouter = createTRPCRouter({
         where: {
           originalLoggedCall: {
             projectId: input.projectId,
-            id: { in: input.selectedLogIds },
+            id: { in: input.loggedCallIds },
           },
           statusCode: 200,
         },

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

@@ -93,17 +93,12 @@ export const promptVariantsRouter = createTRPCRouter({
           visible: true,
         },
       });
-      const outputCount = await prisma.scenarioVariantCell.count({
+      const finishedCount = await prisma.scenarioVariantCell.count({
         where: {
           promptVariantId: input.variantId,
           testScenario: { visible: true },
-          modelResponses: {
+          retrievalStatus: {
-            some: {
+            in: ["COMPLETE", "ERROR"],
-              outdated: false,
-              respPayload: {
-                not: Prisma.AnyNull,
-              },
-            },
           },
         },
       });
@@ -131,7 +126,7 @@ export const promptVariantsRouter = createTRPCRouter({
       const inputTokens = overallTokens._sum?.inputTokens ?? 0;
       const outputTokens = overallTokens._sum?.outputTokens ?? 0;
 
-      const awaitingCompletions = outputCount < scenarioCount;
+      const awaitingCompletions = finishedCount < scenarioCount;
 
       const awaitingEvals = !!evalResults.find(
         (result) => result.totalCount < scenarioCount * evals.length,
@@ -143,7 +138,7 @@ export const promptVariantsRouter = createTRPCRouter({
         outputTokens,
         overallCost: overallTokens._sum?.cost ?? 0,
         scenarioCount,
-        outputCount,
+        finishedCount,
         awaitingCompletions,
         awaitingEvals,
       };

app/src/state/selectedDatasetEntriesSlice.ts

@@ -0,0 +1,33 @@
+import { type SliceCreator } from "./store";
+
+export type SelectedDatasetEntriesSlice = {
+  selectedIds: Set<string>;
+  toggleSelectedId: (id: string) => void;
+  addSelectedIds: (ids: string[]) => void;
+  clearSelectedIds: () => void;
+};
+
+export const createSelectedDatasetEntriesSlice: SliceCreator<SelectedDatasetEntriesSlice> = (
+  set,
+) => ({
+  selectedIds: new Set(),
+  toggleSelectedId: (id: string) =>
+    set((state) => {
+      if (state.selectedDatasetEntries.selectedIds.has(id)) {
+        state.selectedDatasetEntries.selectedIds.delete(id);
+      } else {
+        state.selectedDatasetEntries.selectedIds.add(id);
+      }
+    }),
+  addSelectedIds: (ids: string[]) =>
+    set((state) => {
+      state.selectedDatasetEntries.selectedIds = new Set([
+        ...state.selectedDatasetEntries.selectedIds,
+        ...ids,
+      ]);
+    }),
+  clearSelectedIds: () =>
+    set((state) => {
+      state.selectedDatasetEntries.selectedIds = new Set();
+    }),
+});

app/src/state/selectedLogsSlice.ts

@@ -7,7 +7,7 @@ export type SelectedLogsSlice = {
   clearSelectedLogIds: () => void;
 };
 
-export const createSelectedLogsSlice: SliceCreator<SelectedLogsSlice> = (set, get) => ({
+export const createSelectedLogsSlice: SliceCreator<SelectedLogsSlice> = (set) => ({
   selectedLogIds: new Set(),
   toggleSelectedLogId: (id: string) =>
     set((state) => {

app/src/state/sharedArgumentsEditor.slice.ts

@@ -0,0 +1,33 @@
+import loader, { type Monaco } from "@monaco-editor/loader";
+
+import { type SliceCreator } from "./store";
+
+export const editorBackground = "#fafafa";
+
+export type SharedArgumentsEditorSlice = {
+  monaco: null | Monaco;
+  loadMonaco: () => Promise<void>;
+};
+
+export const createArgumentsEditorSlice: SliceCreator<SharedArgumentsEditorSlice> = (set, get) => ({
+  monaco: loader.__getMonacoInstance(),
+  loadMonaco: async () => {
+    // We only want to run this client-side
+    if (typeof window === "undefined") return;
+
+    const monaco = await loader.init();
+
+    monaco.editor.defineTheme("customTheme", {
+      base: "vs",
+      inherit: true,
+      rules: [],
+      colors: {
+        "editor.background": "#ffffff",
+      },
+    });
+
+    set((state) => {
+      state.sharedArgumentsEditor.monaco = monaco;
+    });
+  },
+});

app/src/state/store.ts

@@ -7,9 +7,17 @@ import {
   type SharedVariantEditorSlice,
   createVariantEditorSlice,
 } from "./sharedVariantEditor.slice";
+import {
+  type SharedArgumentsEditorSlice,
+  createArgumentsEditorSlice,
+} from "./sharedArgumentsEditor.slice";
 import { type APIClient } from "~/utils/api";
 import { type PersistedState, persistOptions } from "./persist";
 import { type SelectedLogsSlice, createSelectedLogsSlice } from "./selectedLogsSlice";
+import {
+  type SelectedDatasetEntriesSlice,
+  createSelectedDatasetEntriesSlice,
+} from "./selectedDatasetEntriesSlice";
 import { type LogFiltersSlice, createLogFiltersSlice } from "./logFiltersSlice";
 import { type ColumnVisibilitySlice, createColumnVisibilitySlice } from "./columnVisiblitySlice";
 import { type FeatureFlagsSlice, createFeatureFlagsSlice } from "./featureFlags";
@@ -18,15 +26,14 @@ enableMapSet();
 
 export type State = {
   isRehydrated: boolean;
-  drawerOpen: boolean;
-  openDrawer: () => void;
-  closeDrawer: () => void;
   api: APIClient | null;
   setApi: (api: APIClient) => void;
   sharedVariantEditor: SharedVariantEditorSlice;
+  sharedArgumentsEditor: SharedArgumentsEditorSlice;
   selectedProjectId: string | null;
   setSelectedProjectId: (id: string) => void;
   selectedLogs: SelectedLogsSlice;
+  selectedDatasetEntries: SelectedDatasetEntriesSlice;
   logFilters: LogFiltersSlice;
   columnVisibility: ColumnVisibilitySlice;
   featureFlags: FeatureFlagsSlice;
@@ -46,22 +53,15 @@ const useBaseStore = create<State, [["zustand/persist", PersistedState], ["zusta
         set((state) => {
           state.api = api;
         }),
-      drawerOpen: false,
-      openDrawer: () =>
-        set((state) => {
-          state.drawerOpen = true;
-        }),
-      closeDrawer: () =>
-        set((state) => {
-          state.drawerOpen = false;
-        }),
       sharedVariantEditor: createVariantEditorSlice(set, get, ...rest),
+      sharedArgumentsEditor: createArgumentsEditorSlice(set, get, ...rest),
       selectedProjectId: null,
       setSelectedProjectId: (id: string) =>
         set((state) => {
           state.selectedProjectId = id;
         }),
       selectedLogs: createSelectedLogsSlice(set, get, ...rest),
+      selectedDatasetEntries: createSelectedDatasetEntriesSlice(set, get, ...rest),
       logFilters: createLogFiltersSlice(set, get, ...rest),
       columnVisibility: createColumnVisibilitySlice(set, get, ...rest),
       featureFlags: createFeatureFlagsSlice(set, get, ...rest),

app/src/utils/countTokens.ts

@@ -12,6 +12,13 @@ export const countOpenAIChatTokens = (
   model: SupportedModel,
   messages: ChatCompletion.Choice.Message[],
 ) => {
-  return new GPTTokens({ model, messages: messages as unknown as GPTTokensMessageItem[] })
+  const reformattedMessages = messages.map((message) => ({
-    .usedTokens;
+    role: message.role,
+    // Not completely accurate, but gives a rough idea of the token count
+    content: message.content ?? JSON.stringify(message.function_call),
+  }));
+  return new GPTTokens({
+    model,
+    messages: reformattedMessages as unknown as GPTTokensMessageItem[],
+  }).usedTokens;
 };

app/src/utils/hooks.ts

@@ -148,6 +148,49 @@ export const useScenarioVars = () => {
   );
 };
 
+export const useDatasets = () => {
+  const selectedProjectId = useAppStore((state) => state.selectedProjectId);
+  return api.datasets.list.useQuery(
+    { projectId: selectedProjectId ?? "" },
+    { enabled: !!selectedProjectId },
+  );
+};
+
+export const useDataset = () => {
+  const router = useRouter();
+  const dataset = api.datasets.get.useQuery(
+    { id: router.query.id as string },
+    { enabled: !!router.query.id },
+  );
+
+  return dataset;
+};
+
+export const useDatasetEntries = () => {
+  const dataset = useDataset().data;
+  const { page, pageSize } = usePageParams();
+
+  const { data, isLoading, ...rest } = api.datasetEntries.list.useQuery(
+    { datasetId: dataset?.id ?? "", page, pageSize },
+    { enabled: !!dataset?.id },
+  );
+
+  const [stableData, setStableData] = useState(data);
+
+  useEffect(() => {
+    // Prevent annoying flashes while logs are loading from the server
+    if (!isLoading) {
+      setStableData(data);
+    }
+  }, [data, isLoading]);
+
+  return { data: stableData, isLoading, ...rest };
+};
+
+export const useDatasetEntry = (entryId: string | null) => {
+  return api.datasetEntries.get.useQuery({ id: entryId as string }, { enabled: !!entryId });
+};
+
 export const useLoggedCalls = (applyFilters = true) => {
   const selectedProjectId = useAppStore((state) => state.selectedProjectId);
   const { page, pageSize } = usePageParams();

app/src/utils/utils.ts

@@ -10,3 +10,45 @@ export const lookupModel = (provider: string, model: string) => {
 
 export const modelLabel = (provider: string, model: string) =>
   `${provider}/${lookupModel(provider, model)?.name ?? model}`;
+
+// Check if the str could be parsed to a message function call
+export const parseableToFunctionCall = (str: string) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  let parsedJSON: any;
+  try {
+    parsedJSON = JSON.parse(str);
+  } catch {
+    return false;
+  }
+
+  // Check if the parsedJSON is an object and not null
+  if (typeof parsedJSON !== "object" || parsedJSON === null) {
+    return false;
+  }
+
+  // Check if only the keys "name" and "arguments" exist
+  const keys = Object.keys(parsedJSON as Record<string, unknown>);
+  if (keys.length !== 2 || !keys.includes("name") || !keys.includes("arguments")) {
+    return false;
+  }
+
+  // Check if both "name" and "arguments" are of type string
+  if (typeof parsedJSON.name !== "string" || typeof parsedJSON.arguments !== "string") {
+    return false;
+  }
+
+  // Check if the "arguments" value is parseable to an object
+  let parsedArguments: unknown;
+  try {
+    parsedArguments = JSON.parse(parsedJSON["arguments"]);
+  } catch {
+    return false;
+  }
+
+  // Check if parsedArguments is an object and not null
+  if (typeof parsedArguments !== "object" || parsedArguments === null) {
+    return false;
+  }
+
+  return true;
+};