feat: Support datasetId in initDataset()

.changeset/support-dataset-id.md

@@ -0,0 +1,5 @@
+---
+"braintrust": minor
+---
+
+feat: Support `datasetId` in `initDataset()`

js/dev/server.ts

@@ -324,14 +324,9 @@ async function getDataset(
       _internal_btql: data._internal_btql ?? undefined,
     });
   } else if ("dataset_id" in data) {
-    const datasetInfo = await getDatasetById({
-      state,
-      datasetId: data.dataset_id,
-    });
     return initDataset({
       state,
-      projectId: datasetInfo.projectId,
-      dataset: datasetInfo.dataset,
+      datasetId: data.dataset_id,
       version: data.dataset_version ?? undefined,
       environment: data.dataset_environment ?? undefined,
       _internal_btql: data._internal_btql ?? undefined,
@@ -344,27 +339,6 @@ async function getDataset(
   }
 }
 
-const datasetFetchSchema = z.object({
-  project_id: z.string(),
-  name: z.string(),
-});
-async function getDatasetById({
-  state,
-  datasetId,
-}: {
-  state: BraintrustState;
-  datasetId: string;
-}): Promise<{ projectId: string; dataset: string }> {
-  const dataset = await state.appConn().post_json("api/dataset/get", {
-    id: datasetId,
-  });
-  const parsed = z.array(datasetFetchSchema).parse(dataset);
-  if (parsed.length === 0) {
-    throw new Error(`Dataset '${datasetId}' not found`);
-  }
-  return { projectId: parsed[0].project_id, dataset: parsed[0].name };
-}
-
 function makeScorer(
   state: BraintrustState,
   name: string,

js/src/logger.test.ts

@@ -451,6 +451,180 @@ test("init validation", () => {
   );
 });
 
+test("initDataset supports dataset IDs without a project", async () => {
+  const state = await _exportsForTestingOnly.simulateLoginForTests();
+
+  try {
+    vi.spyOn(state, "login").mockResolvedValue(state as any);
+    const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue([
+      {
+        object_id: "00000000-0000-0000-0000-000000000002",
+        object_name: "test-dataset",
+        parent_cols: {
+          project: {
+            id: "00000000-0000-0000-0000-000000000001",
+            name: "test-project",
+          },
+        },
+      },
+    ]);
+
+    const datasetById = initDataset({
+      datasetId: "00000000-0000-0000-0000-000000000002",
+      state,
+    });
+    await expect(datasetById.id).resolves.toBe(
+      "00000000-0000-0000-0000-000000000002",
+    );
+    await expect(datasetById.name).resolves.toBe("test-dataset");
+    await expect(datasetById.project).resolves.toMatchObject({
+      id: "00000000-0000-0000-0000-000000000001",
+      name: "test-project",
+    });
+
+    expect(postJson).toHaveBeenCalledOnce();
+    expect(postJson).toHaveBeenCalledWith("api/self/get_object_info", {
+      object_type: "dataset",
+      object_ids: ["00000000-0000-0000-0000-000000000002"],
+    });
+  } finally {
+    _exportsForTestingOnly.simulateLogoutForTests();
+    vi.restoreAllMocks();
+  }
+});
+
+test("initDataset gives dataset IDs precedence over names", async () => {
+  const state = await _exportsForTestingOnly.simulateLoginForTests();
+
+  try {
+    vi.spyOn(state, "login").mockResolvedValue(state as any);
+    const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue([
+      {
+        object_id: "00000000-0000-0000-0000-000000000002",
+        object_name: "test-dataset",
+        parent_cols: {
+          project: {
+            id: "00000000-0000-0000-0000-000000000001",
+            name: "test-project",
+          },
+        },
+      },
+    ]);
+
+    const datasetByIdAndName = initDataset({
+      project: "ignored-project",
+      dataset: "ignored-dataset",
+      datasetId: "00000000-0000-0000-0000-000000000002",
+      state,
+    });
+    await datasetByIdAndName.id;
+
+    expect(postJson).toHaveBeenCalledOnce();
+    expect(postJson).toHaveBeenCalledWith("api/self/get_object_info", {
+      object_type: "dataset",
+      object_ids: ["00000000-0000-0000-0000-000000000002"],
+    });
+    expect(postJson).not.toHaveBeenCalledWith(
+      "api/dataset/register",
+      expect.anything(),
+    );
+  } finally {
+    _exportsForTestingOnly.simulateLogoutForTests();
+    vi.restoreAllMocks();
+  }
+});
+
+test("initDataset keeps the existing name-based registration path", async () => {
+  const state = await _exportsForTestingOnly.simulateLoginForTests();
+
+  try {
+    vi.spyOn(state, "login").mockResolvedValue(state as any);
+    const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue({
+      project: {
+        id: "00000000-0000-0000-0000-000000000001",
+        name: "test-project",
+      },
+      dataset: {
+        id: "00000000-0000-0000-0000-000000000002",
+        name: "test-dataset",
+      },
+    });
+
+    const datasetByName = initDataset({
+      project: "test-project",
+      dataset: "test-dataset",
+      state,
+    });
+    await datasetByName.id;
+
+    expect(postJson).toHaveBeenCalledOnce();
+    expect(postJson).toHaveBeenCalledWith(
+      "api/dataset/register",
+      expect.objectContaining({
+        project_name: "test-project",
+        dataset_name: "test-dataset",
+      }),
+    );
+  } finally {
+    _exportsForTestingOnly.simulateLogoutForTests();
+    vi.restoreAllMocks();
+  }
+});
+
+test("initDataset reports an unknown dataset ID", async () => {
+  const state = await _exportsForTestingOnly.simulateLoginForTests();
+
+  try {
+    vi.spyOn(state, "login").mockResolvedValue(state as any);
+    const postJson = vi
+      .spyOn(state.appConn(), "post_json")
+      .mockResolvedValue([]);
+
+    const dataset = initDataset({
+      datasetId: "00000000-0000-0000-0000-000000000099",
+      state,
+    });
+
+    await expect(dataset.id).rejects.toThrow(
+      "Dataset with ID 00000000-0000-0000-0000-000000000099 not found",
+    );
+    expect(postJson).toHaveBeenCalledOnce();
+    expect(postJson).toHaveBeenCalledWith("api/self/get_object_info", {
+      object_type: "dataset",
+      object_ids: ["00000000-0000-0000-0000-000000000099"],
+    });
+    expect(postJson).not.toHaveBeenCalledWith(
+      "api/dataset/register",
+      expect.anything(),
+    );
+  } finally {
+    _exportsForTestingOnly.simulateLogoutForTests();
+    vi.restoreAllMocks();
+  }
+});
+
+test("initDataset rejects dataset ID updates", () => {
+  const optionsWithDescription = {
+    datasetId: "00000000-0000-0000-0000-000000000002",
+    description: "unsupported update",
+  };
+  const optionsWithMetadata = {
+    datasetId: "00000000-0000-0000-0000-000000000002",
+    metadata: { unsupported: "update" },
+  };
+
+  expect(() => {
+    initDataset(optionsWithDescription);
+  }).toThrow(
+    "Cannot specify description or metadata when datasetId is provided",
+  );
+  expect(() => {
+    initDataset(optionsWithMetadata);
+  }).toThrow(
+    "Cannot specify description or metadata when datasetId is provided",
+  );
+});
+
 test("init accepts dataset with id only", () => {
   // Test that the type system accepts {id: string}
   const datasetIdOnly = { id: "dataset-id-123" };

js/src/logger.ts

@@ -116,6 +116,17 @@ const datasetSnapshotRegisterResponseSchema = z.object({
   found_existing: z.boolean().optional(),
 });
 
+const datasetObjectInfoSchema = z.object({
+  object_id: z.string(),
+  object_name: z.string(),
+  parent_cols: z.object({
+    project: z.object({
+      id: z.string(),
+      name: z.string(),
+    }),
+  }),
+});
+
 const datasetRestorePreviewResultSchema = z.object({
   rows_to_restore: z.number(),
   rows_to_delete: z.number(),
@@ -4469,6 +4480,7 @@ declare global {
 export type InitDatasetOptions<IsLegacyDataset extends boolean> =
   FullLoginOptions & {
     dataset?: string;
+    datasetId?: string;
     description?: string;
     version?: string;
     environment?: string;
@@ -4690,22 +4702,23 @@ async function serializeDatasetForExperiment({
 }
 
 /**
- * Create a new dataset in a specified project. If the project does not exist, it will be created.
+ * Initialize a dataset by name or ID. When initializing by name, the dataset and its project will be created if they do not exist.
  *
  * @param options Options for configuring initDataset().
- * @param options.project The name of the project to create the dataset in. Must specify at least one of `project` or `projectId`.
- * @param options.dataset The name of the dataset to create. If not specified, a name will be generated automatically.
- * @param options.description An optional description of the dataset.
+ * @param options.project The name of the project to create the dataset in. Must specify at least one of `project` or `projectId` unless `datasetId` is provided.
+ * @param options.dataset The name of the dataset to create. If not specified, a name will be generated automatically. Ignored if `datasetId` is provided.
+ * @param options.datasetId The ID of an existing dataset. Takes precedence over `dataset` and does not require `project` or `projectId`.
+ * @param options.description An optional description when initializing a dataset by name. Cannot be used with `datasetId`.
  * @param options.version Pin the dataset to a specific version xact_id. If `snapshotName` or `environment` are also provided, `version` takes precedence.
  * @param options.snapshotName Pin the dataset to the version captured by this named snapshot. If `environment` is also provided, `snapshotName` takes precedence.
  * @param options.environment Pin the dataset to the version tagged with this environment slug.
  * @param options.appUrl The URL of the Braintrust App. Defaults to https://www.braintrust.dev.
  * @param options.apiKey The API key to use. If the parameter is not specified, will try to use the `BRAINTRUST_API_KEY` environment variable. In Node.js, if that is unset, will try the nearest `.env.braintrust` file in the current working directory or parent directories. If no API key is specified, will prompt the user to login.
  * @param options.orgName (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple.
- * @param options.projectId The id of the project to create the dataset in. This takes precedence over `project` if specified.
- * @param options.metadata A dictionary with additional data about the dataset. The values in `metadata` can be any JSON-serializable type, but its keys must be strings.
+ * @param options.projectId The id of the project to create the dataset in. This takes precedence over `project` if specified and is not required when `datasetId` is provided.
+ * @param options.metadata A dictionary with additional data when initializing a dataset by name. Cannot be used with `datasetId`. The values in `metadata` can be any JSON-serializable type, but its keys must be strings.
  * @param options.useOutput (Deprecated) If true, records will be fetched from this dataset in the legacy format, with the "expected" field renamed to "output". This option will be removed in a future version of Braintrust.
- * @returns The newly created Dataset.
+ * @returns The initialized Dataset.
  */
 export function initDataset<
   IsLegacyDataset extends boolean = typeof DEFAULT_IS_LEGACY_DATASET,
@@ -4753,6 +4766,7 @@ export function initDataset<
   const {
     project,
     dataset,
+    datasetId,
     description,
     version,
     snapshotName,
@@ -4768,6 +4782,15 @@ export function initDataset<
     state: stateArg,
     _internal_btql,
   } = options;
+  if (
+    datasetId !== undefined &&
+    (description !== undefined || metadata !== undefined)
+  ) {
+    throw new Error(
+      "Cannot specify description or metadata when datasetId is provided",
+    );
+  }
+
   const selection = normalizeDatasetSelection({
     version,
     environment,
@@ -4799,6 +4822,37 @@ export function initDataset<
         forceLogin,
       });
 
+      if (datasetId !== undefined) {
+        const objectInfo = datasetObjectInfoSchema.array().parse(
+          await state.appConn().post_json("api/self/get_object_info", {
+            object_type: "dataset",
+            object_ids: [datasetId],
+          }),
+        );
+        if (objectInfo.length === 0) {
+          throw new Error(`Dataset with ID ${datasetId} not found`);
+        }
+        if (objectInfo.length !== 1) {
+          throw new Error(
+            `Expected exactly one dataset with ID ${datasetId}, but found ${objectInfo.length}`,
+          );
+        }
+
+        const datasetInfo = objectInfo[0];
+        return {
+          project: {
+            id: datasetInfo.parent_cols.project.id,
+            name: datasetInfo.parent_cols.project.name,
+            fullInfo: datasetInfo.parent_cols.project,
+          },
+          dataset: {
+            id: datasetInfo.object_id,
+            name: datasetInfo.object_name,
+            fullInfo: datasetInfo,
+          },
+        };
+      }
+
       const args: Record<string, unknown> = {
         org_id: state.orgId,
         project_name: project,

js/src/wrappers/vitest/README.md

@@ -239,13 +239,34 @@ bt.describe("Translation Suite", () => {
 import { initDataset } from "braintrust";
 
 initDataset({
-  project: "my-project", // Project name or ID
-  dataset: "my-dataset", // Dataset name or ID
+  project: "my-project", // Project name
+  dataset: "my-dataset", // Dataset name
   version: "v1.2", // Optional: specific version
   description: "...", // Optional: description
 }).fetchedData();
 ```
 
+Use `projectId` instead of `project` to select the project by ID while
+initializing a dataset by name:
+
+```typescript
+initDataset({
+  projectId: "my-project-id",
+  dataset: "my-dataset",
+}).fetchedData();
+```
+
+Use `datasetId` to load an existing dataset directly. A project is not
+required, and `datasetId` takes precedence if project or dataset names are also
+provided:
+
+```typescript
+initDataset({
+  datasetId: "my-dataset-id",
+  version: "v1.2",
+}).fetchedData();
+```
+
 ### Combining Datasets with Scorers
 
 The most powerful pattern combines datasets with automatic scoring: