Adapt batch evaluations for datasets v2 (#946)

Author: andresgutgon

apps/web/src/actions/documents/runDocumentInBatchAction.ts

@@ -32,6 +32,7 @@ export const runDocumentInBatchAction = withDataset
                   prompt: ctx.document.content,
                 })
           const docParams = metadata.parameters
+          // @ts-expect-error - dataset V2 are not send here yet
           const headers = ctx.dataset.fileMetadata.headers
           const paramKeys = Object.keys(parameters)
           Array.from(docParams).forEach((key) => {

apps/web/src/actions/evaluations/_helpers.ts

@@ -1,18 +1,38 @@
-import { DatasetsRepository } from '@latitude-data/core/repositories'
+import {
+  DatasetsRepository,
+  DatasetsV2Repository,
+} from '@latitude-data/core/repositories'
 import { z } from 'zod'
 import { createServerActionProcedure } from 'zsa'
 
 import { withDocument } from '../procedures'
+import { DatasetVersion } from '@latitude-data/constants'
+import { Dataset, DatasetV2 } from '@latitude-data/core/browser'
 
 export const withDataset = createServerActionProcedure(withDocument)
-  .input(z.object({ datasetId: z.number() }))
+  .input(
+    z.object({
+      datasetId: z.number(),
+      datasetVersion: z.nativeEnum(DatasetVersion),
+    }),
+  )
   .handler(async ({ input, ctx }) => {
-    const datasetsRepo = new DatasetsRepository(ctx.workspace.id)
-    const dataset = await datasetsRepo
-      .find(input.datasetId)
-      .then((r) => r.unwrap())
+    const datasetVersion = input.datasetVersion
+    let response = { ...ctx, datasetVersion }
 
-    return { ...ctx, dataset }
+    // DEPRECATED
+    if (datasetVersion === DatasetVersion.V1) {
+      const datasetsRepo = new DatasetsRepository(ctx.workspace.id)
+      const dataset = await datasetsRepo
+        .find(input.datasetId)
+        .then((r) => r.unwrap())
+      return { ...response, dataset: dataset as Dataset }
+    }
+
+    const repo = new DatasetsV2Repository(ctx.workspace.id)
+    const dataset = await repo.find(input.datasetId).then((r) => r.unwrap())
+
+    return { ...response, dataset: dataset as DatasetV2 }
   })
 
 export const USER_DECIDED_TO_IGNORE_THIS_PARAMETER = -1