From 4e993e881712e21c9024a5ec8444e2ba8326907f Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:34:51 +0000 Subject: [PATCH 1/3] feat: Support `datasetId` in `initDataset()` --- js/dev/server.ts | 28 +---- js/src/logger.test.ts | 176 +++++++++++++++++++++++++++++++ js/src/logger.ts | 92 ++++++++++++++-- js/src/wrappers/vitest/README.md | 25 ++++- 4 files changed, 281 insertions(+), 40 deletions(-) diff --git a/js/dev/server.ts b/js/dev/server.ts index 5db8b467f..3118e8252 100644 --- a/js/dev/server.ts +++ b/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, diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index a061c8399..44aae9fbc 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -451,6 +451,182 @@ 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(() => { + // @ts-expect-error description cannot be used with datasetId + initDataset(optionsWithDescription); + }).toThrow( + "Cannot specify description or metadata when datasetId is provided", + ); + expect(() => { + // @ts-expect-error metadata cannot be used with datasetId + 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" }; diff --git a/js/src/logger.ts b/js/src/logger.ts index ad316c481..5dc26866b 100644 --- a/js/src/logger.ts +++ b/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(), @@ -4466,19 +4477,36 @@ declare global { var __bt_eval_internal_btql: Record | undefined; } -export type InitDatasetOptions = +type InitDatasetBaseOptions = FullLoginOptions & { - dataset?: string; - description?: string; version?: string; environment?: string; snapshotName?: string; projectId?: string; - metadata?: Record; state?: BraintrustState; _internal_btql?: Record; } & UseOutputOption; +type InitDatasetByIdOptions = + InitDatasetBaseOptions & { + datasetId: string; + dataset?: string; + description?: never; + metadata?: never; + }; + +type InitDatasetByNameOptions = + InitDatasetBaseOptions & { + datasetId?: never; + dataset?: string; + description?: string; + metadata?: Record; + }; + +export type InitDatasetOptions = + | InitDatasetByIdOptions + | InitDatasetByNameOptions; + export type FullInitDatasetOptions = { project?: string; } & InitDatasetOptions; @@ -4690,22 +4718,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 +4782,7 @@ export function initDataset< const { project, dataset, + datasetId, description, version, snapshotName, @@ -4768,6 +4798,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 +4838,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 = { org_id: state.orgId, project_name: project, diff --git a/js/src/wrappers/vitest/README.md b/js/src/wrappers/vitest/README.md index c2ddb66df..ef5885ee3 100644 --- a/js/src/wrappers/vitest/README.md +++ b/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: From 7b3d947cc22a3f8e5d653bdff68f162d71da9c80 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:43:03 +0000 Subject: [PATCH 2/3] Update PR #2398 --- .changeset/support-dataset-id.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/support-dataset-id.md diff --git a/.changeset/support-dataset-id.md b/.changeset/support-dataset-id.md new file mode 100644 index 000000000..9b8808be0 --- /dev/null +++ b/.changeset/support-dataset-id.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +feat: Support `datasetId` in `initDataset()` From abdff1972bdd0bad25b90a27d0609c675d181975 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:31:36 +0000 Subject: [PATCH 3/3] Update PR #2398 --- .changeset/support-dataset-id.md | 2 +- js/src/logger.test.ts | 2 -- js/src/logger.ts | 26 +++++--------------------- 3 files changed, 6 insertions(+), 24 deletions(-) diff --git a/.changeset/support-dataset-id.md b/.changeset/support-dataset-id.md index 9b8808be0..60748aeaa 100644 --- a/.changeset/support-dataset-id.md +++ b/.changeset/support-dataset-id.md @@ -1,5 +1,5 @@ --- -"braintrust": patch +"braintrust": minor --- feat: Support `datasetId` in `initDataset()` diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index 44aae9fbc..ddb2a8408 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -614,13 +614,11 @@ test("initDataset rejects dataset ID updates", () => { }; expect(() => { - // @ts-expect-error description cannot be used with datasetId initDataset(optionsWithDescription); }).toThrow( "Cannot specify description or metadata when datasetId is provided", ); expect(() => { - // @ts-expect-error metadata cannot be used with datasetId initDataset(optionsWithMetadata); }).toThrow( "Cannot specify description or metadata when datasetId is provided", diff --git a/js/src/logger.ts b/js/src/logger.ts index 5dc26866b..7457ddbe5 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -4477,36 +4477,20 @@ declare global { var __bt_eval_internal_btql: Record | undefined; } -type InitDatasetBaseOptions = +export type InitDatasetOptions = FullLoginOptions & { + dataset?: string; + datasetId?: string; + description?: string; version?: string; environment?: string; snapshotName?: string; projectId?: string; + metadata?: Record; state?: BraintrustState; _internal_btql?: Record; } & UseOutputOption; -type InitDatasetByIdOptions = - InitDatasetBaseOptions & { - datasetId: string; - dataset?: string; - description?: never; - metadata?: never; - }; - -type InitDatasetByNameOptions = - InitDatasetBaseOptions & { - datasetId?: never; - dataset?: string; - description?: string; - metadata?: Record; - }; - -export type InitDatasetOptions = - | InitDatasetByIdOptions - | InitDatasetByNameOptions; - export type FullInitDatasetOptions = { project?: string; } & InitDatasetOptions;