diff --git a/.changeset/goofy-hotels-care.md b/.changeset/goofy-hotels-care.md new file mode 100644 index 000000000..ede100873 --- /dev/null +++ b/.changeset/goofy-hotels-care.md @@ -0,0 +1,69 @@ +--- +"braintrust": minor +--- + +Add dataset versioning support to `init()`, `initDataset()`, and dataset objects. + +You can now pin dataset reads and experiment registration by explicit version, snapshot name, or environment tag: + +```ts +import { init, initDataset } from "braintrust"; + +const datasetByVersion = initDataset({ + project: "support-bot", + dataset: "production-cases", + version: "1234567890123456", +}); + +const datasetBySnapshot = initDataset({ + project: "support-bot", + dataset: "production-cases", + snapshotName: "baseline", +}); + +const datasetByEnvironment = initDataset({ + project: "support-bot", + dataset: "production-cases", + environment: "production", +}); + +init({ + project: "support-bot", + experiment: "baseline-eval", + dataset: { + id: "00000000-0000-0000-0000-000000000123", + snapshotName: "baseline", + }, +}); +``` + +Dataset objects now expose snapshot CRUD helpers, plus lookup by snapshot name or xact id: + +```ts +const dataset = initDataset({ + project: "support-bot", + dataset: "production-cases", +}); + +const snapshot = await dataset.createSnapshot({ + name: "baseline", + description: "Before the prompt rollout", +}); + +await dataset.updateSnapshot(snapshot.id, { + name: "baseline-v2", + description: null, +}); + +const snapshots = await dataset.listSnapshots(); +const byName = await dataset.getSnapshot({ + snapshotName: "baseline-v2", +}); +const byXactId = await dataset.getSnapshot({ + xactId: snapshot.xact_id, +}); + +await dataset.deleteSnapshot(snapshot.id); +``` + +`braintrust/dev` now also respects `dataset_version` and `dataset_environment` when resolving datasets for evals, so local eval runs match the pinned dataset selection used by the main SDK. diff --git a/js/dev/server.ts b/js/dev/server.ts index 5499fe208..ebce5c3f5 100644 --- a/js/dev/server.ts +++ b/js/dev/server.ts @@ -318,6 +318,8 @@ async function getDataset( state, project: data.project_name, dataset: data.dataset_name, + version: data.dataset_version ?? undefined, + environment: data.dataset_environment ?? undefined, _internal_btql: data._internal_btql ?? undefined, }); } else if ("dataset_id" in data) { @@ -329,6 +331,8 @@ async function getDataset( state, projectId: datasetInfo.projectId, dataset: datasetInfo.dataset, + version: data.dataset_version ?? undefined, + environment: data.dataset_environment ?? undefined, _internal_btql: data._internal_btql ?? undefined, }); } else { diff --git a/js/src/exports.ts b/js/src/exports.ts index f2d940d99..93a12bb05 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -8,6 +8,9 @@ export type { CompiledPromptParams, CompletionPrompt, ContextParentSpanIds, + DatasetRestorePreviewResult, + DatasetRestoreResult, + DatasetSnapshot, DataSummary, DatasetSummary, DefaultMetadataType, diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index fc4f14f7b..c4536bb78 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -4,6 +4,7 @@ import { vi, expect, test, describe, beforeEach, afterEach } from "vitest"; import { _exportsForTestingOnly, init, + initDataset, initLogger, Prompt, BraintrustState, @@ -453,6 +454,804 @@ test("init accepts dataset with id and version", () => { expect(datasetWithVersion.version).toBe("v2"); }); +test("init accepts dataset with id and environment", () => { + const datasetWithEnvironment = { + id: "dataset-id-123", + environment: "production", + }; + + expect(datasetWithEnvironment.id).toBe("dataset-id-123"); + expect(datasetWithEnvironment.environment).toBe("production"); +}); + +test("init accepts dataset with id and snapshotName", () => { + const datasetWithSnapshot = { + id: "dataset-id-123", + snapshotName: "123", + }; + + expect(datasetWithSnapshot.id).toBe("dataset-id-123"); + expect(datasetWithSnapshot.snapshotName).toBe("123"); +}); + +function mockInitGitMetadata() { + vi.spyOn(_exportsForTestingOnly.isomorph, "getRepoInfo").mockResolvedValue( + undefined, + ); + vi.spyOn( + _exportsForTestingOnly.isomorph, + "getPastNAncestors", + ).mockResolvedValue([]); +} + +test("initDataset prefers version over environment in eval data", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + 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 dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + version: "123", + environment: "production", + state, + }); + + await expect(dataset.toEvalData()).resolves.toEqual({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "123", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.toEvalData preserves dataset_environment", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + vi.spyOn(state.apiConn(), "get_json").mockResolvedValue({ + object_version: "123", + }); + 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 dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + environment: "production", + state, + }); + + await expect(dataset.toEvalData()).resolves.toEqual({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_environment: "production", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.toEvalData preserves dataset_snapshot_name", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce([ + { + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "123", + description: null, + xact_id: "456", + created: "2026-03-31T00:00:00.000Z", + }, + ]); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + snapshotName: "123", + state, + }); + + await expect(dataset.toEvalData()).resolves.toEqual({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_snapshot_name: "123", + }); + expect(postJson).toHaveBeenNthCalledWith(2, "api/dataset_snapshot/get", { + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "123", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.version preserves pinned-version fast path", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + const login = vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi.spyOn(state.appConn(), "post_json"); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + version: "123", + state, + }); + + await expect(dataset.version()).resolves.toBe("123"); + expect(login).not.toHaveBeenCalled(); + expect(postJson).not.toHaveBeenCalled(); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.createSnapshot forwards update when requested", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce({ + dataset_snapshot: { + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "snapshot", + description: "updated description", + xact_id: "123", + created: "2026-03-31T00:00:00.000Z", + }, + found_existing: true, + }); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + version: "123", + state, + }); + + await expect( + dataset.createSnapshot({ + name: "snapshot", + description: "updated description", + update: true, + }), + ).resolves.toMatchObject({ + id: "00000000-0000-0000-0000-000000000004", + xact_id: "123", + }); + + expect(postJson).toHaveBeenNthCalledWith(2, "api/dataset_snapshot/register", { + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_snapshot_name: "snapshot", + description: "updated description", + xact_id: "123", + update: true, + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.getSnapshot looks up snapshots by name", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce([ + { + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "snapshot", + description: null, + xact_id: "123", + created: "2026-03-31T00:00:00.000Z", + }, + ]); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + + await expect( + dataset.getSnapshot({ + snapshotName: "snapshot", + }), + ).resolves.toMatchObject({ + id: "00000000-0000-0000-0000-000000000004", + name: "snapshot", + xact_id: "123", + }); + + expect(postJson).toHaveBeenNthCalledWith(2, "api/dataset_snapshot/get", { + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "snapshot", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.getSnapshot looks up snapshots by xact id", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce([ + { + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "snapshot", + description: null, + xact_id: "123", + created: "2026-03-31T00:00:00.000Z", + }, + ]); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + + await expect( + dataset.getSnapshot({ + xactId: "123", + }), + ).resolves.toMatchObject({ + id: "00000000-0000-0000-0000-000000000004", + name: "snapshot", + xact_id: "123", + }); + + expect(postJson).toHaveBeenNthCalledWith(2, "api/dataset_snapshot/get", { + dataset_id: "00000000-0000-0000-0000-000000000002", + xact_id: "123", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.updateSnapshot patches snapshot metadata by id", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce({ + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "renamed snapshot", + description: null, + xact_id: "123", + created: "2026-03-31T00:00:00.000Z", + }); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + + await expect( + dataset.updateSnapshot("00000000-0000-0000-0000-000000000004", { + name: "renamed snapshot", + description: null, + }), + ).resolves.toMatchObject({ + id: "00000000-0000-0000-0000-000000000004", + name: "renamed snapshot", + description: null, + }); + + expect(postJson).toHaveBeenNthCalledWith(2, "api/dataset_snapshot/patch_id", { + id: "00000000-0000-0000-0000-000000000004", + name: "renamed snapshot", + description: null, + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.restorePreview posts restore preview request", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + vi.spyOn(state.appConn(), "post_json").mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }); + const postJson = vi + .spyOn(state.apiConn(), "post_json") + .mockResolvedValueOnce({ + rows_to_restore: 3, + rows_to_delete: 1, + }); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + + await expect( + dataset.restorePreview({ + version: "123", + }), + ).resolves.toEqual({ + rows_to_restore: 3, + rows_to_delete: 1, + }); + + expect(postJson).toHaveBeenNthCalledWith( + 1, + "v1/dataset/00000000-0000-0000-0000-000000000002/restore/preview", + { + version: "123", + }, + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("dataset.restore posts restore request", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + vi.spyOn(state.appConn(), "post_json").mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }); + const postJson = vi + .spyOn(state.apiConn(), "post_json") + .mockResolvedValueOnce({ + xact_id: "456", + rows_restored: 3, + rows_deleted: 1, + }); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + + await expect( + dataset.restore({ + version: "123", + }), + ).resolves.toEqual({ + xact_id: "456", + rows_restored: 3, + rows_deleted: 1, + }); + + expect(postJson).toHaveBeenNthCalledWith( + 1, + "v1/dataset/00000000-0000-0000-0000-000000000002/restore", + { + version: "123", + }, + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init keeps plain dataset refs attached to the experiment", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + vi.spyOn(state.appConn(), "post_json").mockResolvedValue({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + }, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(experiment.dataset).toMatchObject({ + id: "00000000-0000-0000-0000-000000000002", + }); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init resolves dataset version from Dataset instances before experiment registration", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + dataset: { + id: "00000000-0000-0000-0000-000000000002", + name: "test-dataset", + }, + }) + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const dataset = initDataset({ + project: "test-project", + dataset: "test-dataset", + state, + }); + const version = vi.spyOn(dataset, "version").mockResolvedValue("123"); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(version).toHaveBeenCalled(); + expect(postJson).toHaveBeenNthCalledWith( + 2, + "api/experiment/register", + expect.objectContaining({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "123", + }), + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init resolves dataset environment before experiment registration", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + const getJson = vi.spyOn(state.apiConn(), "get_json").mockResolvedValue({ + object_version: "123", + }); + const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + environment: "production", + }, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(getJson).toHaveBeenCalledWith( + "environment-object/dataset/00000000-0000-0000-0000-000000000002/production", + { + org_name: "test-org-name", + }, + ); + expect(experiment.dataset).toMatchObject({ + id: "00000000-0000-0000-0000-000000000002", + environment: "production", + }); + expect(postJson).toHaveBeenCalledWith( + "api/experiment/register", + expect.objectContaining({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "123", + }), + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init resolves dataset environment without org_name when orgName is unset", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + state.orgName = null; + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + const getJson = vi.spyOn(state.apiConn(), "get_json").mockResolvedValue({ + object_version: "123", + }); + const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + environment: "production", + }, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(getJson).toHaveBeenCalledWith( + "environment-object/dataset/00000000-0000-0000-0000-000000000002/production", + ); + expect(postJson).toHaveBeenCalledWith( + "api/experiment/register", + expect.objectContaining({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "123", + }), + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init prefers dataset version over environment before experiment registration", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + const getJson = vi.spyOn(state.apiConn(), "get_json"); + const postJson = vi.spyOn(state.appConn(), "post_json").mockResolvedValue({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + version: "123", + environment: "production", + }, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(getJson).not.toHaveBeenCalled(); + expect(postJson).toHaveBeenCalledWith( + "api/experiment/register", + expect.objectContaining({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "123", + }), + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init resolves dataset snapshots before experiment registration", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + const postJson = vi + .spyOn(state.appConn(), "post_json") + .mockResolvedValueOnce([ + { + id: "00000000-0000-0000-0000-000000000004", + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "123", + description: null, + xact_id: "456", + created: "2026-03-31T00:00:00.000Z", + }, + ]) + .mockResolvedValueOnce({ + project: { + id: "00000000-0000-0000-0000-000000000001", + name: "test-project", + }, + experiment: { + id: "00000000-0000-0000-0000-000000000003", + project_id: "00000000-0000-0000-0000-000000000001", + name: "test-experiment", + public: false, + }, + }); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + snapshotName: "123", + }, + setCurrent: false, + state, + }); + + await experiment.id; + + expect(postJson).toHaveBeenNthCalledWith(1, "api/dataset_snapshot/get", { + dataset_id: "00000000-0000-0000-0000-000000000002", + name: "123", + }); + expect(postJson).toHaveBeenNthCalledWith( + 2, + "api/experiment/register", + expect.objectContaining({ + dataset_id: "00000000-0000-0000-0000-000000000002", + dataset_version: "456", + }), + ); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + +test("init surfaces dataset environment lookup errors instead of falling back to latest", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state); + mockInitGitMetadata(); + vi.spyOn(state.apiConn(), "get_json").mockRejectedValue( + new Error("environment lookup failed"), + ); + const postJson = vi.spyOn(state.appConn(), "post_json"); + + const experiment = init({ + project: "test-project", + experiment: "test-experiment", + dataset: { + id: "00000000-0000-0000-0000-000000000002", + environment: "production", + }, + setCurrent: false, + state, + }); + + await expect(experiment.id).rejects.toThrow("environment lookup failed"); + expect(postJson).not.toHaveBeenCalled(); + + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + describe("loader version precedence", () => { let state: BraintrustState; let getJson: ReturnType; diff --git a/js/src/logger.ts b/js/src/logger.ts index 2dd3d5150..ef8ed3a65 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -69,6 +69,8 @@ import { type GitMetadataSettingsType as GitMetadataSettings, type ChatCompletionMessageParamType as Message, type ChatCompletionOpenAIMessageParamType as OpenAIMessage, + DatasetSnapshot as datasetSnapshotSchema, + type DatasetSnapshotType as DatasetSnapshot, PromptData as promptDataSchema, type PromptDataType as PromptData, Prompt as promptSchema, @@ -90,6 +92,28 @@ const RESET_CONTEXT_MANAGER_STATE = Symbol.for( // 6 MB for the AWS lambda gateway (from our own testing). export const DEFAULT_MAX_REQUEST_SIZE = 6 * 1024 * 1024; +export type { DatasetSnapshot }; + +const datasetSnapshotRegisterResponseSchema = z.object({ + dataset_snapshot: datasetSnapshotSchema, + found_existing: z.boolean().optional(), +}); + +const datasetRestorePreviewResultSchema = z.object({ + rows_to_restore: z.number(), + rows_to_delete: z.number(), +}); +export type DatasetRestorePreviewResult = z.infer< + typeof datasetRestorePreviewResultSchema +>; + +const datasetRestoreResultSchema = z.object({ + xact_id: z.string().nullable(), + rows_restored: z.number(), + rows_deleted: z.number(), +}); +export type DatasetRestoreResult = z.infer; + const parametersRowSchema = z.object({ id: z.string().uuid(), _xact_id: z.string(), @@ -3390,13 +3414,55 @@ type InitOpenOption = { }; /** - * Reference to a dataset by ID and optional version. + * Reference to a dataset by ID and optional explicit selector. */ -export interface DatasetRef { - id: string; +type DatasetSelection = { version?: string; + environment?: string; + snapshotName?: string; +}; + +type DatasetSnapshotNameLookup = { + snapshotName: string; + xactId?: never; +}; + +type DatasetSnapshotXactLookup = { + snapshotName?: never; + xactId: string; +}; + +export type DatasetSnapshotLookup = + | DatasetSnapshotNameLookup + | DatasetSnapshotXactLookup; + +function isDatasetSnapshotNameLookup( + lookup: DatasetSnapshotLookup, +): lookup is DatasetSnapshotNameLookup { + return "snapshotName" in lookup; } +function assertDatasetSnapshotLookup(lookup: { + snapshotName?: string; + xactId?: string; +}): asserts lookup is DatasetSnapshotLookup { + const hasSnapshotName = lookup.snapshotName !== undefined; + const hasXactId = lookup.xactId !== undefined; + if (hasSnapshotName === hasXactId) { + throw new Error("Exactly one of snapshotName or xactId must be provided"); + } +} + +type DatasetPinState = { + lazyPinnedVersion?: LazyValue; + pinnedEnvironment?: string; + pinnedSnapshotName?: string; +}; + +export type DatasetRef = { + id: string; +} & DatasetSelection; + export interface ParametersRef { id: string; version?: string; @@ -3621,20 +3687,13 @@ export function init( } if (dataset !== undefined) { - if ( - "id" in dataset && - typeof dataset.id === "string" && - !("__braintrust_dataset_marker" in dataset) - ) { - // Simple {id: ..., version?: ...} object - args["dataset_id"] = dataset.id; - if ("version" in dataset && dataset.version !== undefined) { - args["dataset_version"] = dataset.version; - } - } else { - // Full Dataset object - args["dataset_id"] = await (dataset as AnyDataset).id; - args["dataset_version"] = await (dataset as AnyDataset).version(); + const datasetSelection = await serializeDatasetForExperiment({ + dataset, + state, + }); + args["dataset_id"] = datasetSelection.datasetId; + if (datasetSelection.datasetVersion !== undefined) { + args["dataset_version"] = datasetSelection.datasetVersion; } } @@ -3704,9 +3763,7 @@ export function init( const ret = new Experiment( state, lazyMetadata, - dataset !== undefined && "version" in dataset - ? (dataset as AnyDataset) - : undefined, + dataset !== undefined ? (dataset as AnyDataset) : undefined, ); if (options.setCurrent ?? true) { state.currentExperiment = ret; @@ -3795,6 +3852,8 @@ export type InitDatasetOptions = dataset?: string; description?: string; version?: string; + environment?: string; + snapshotName?: string; projectId?: string; metadata?: Record; state?: BraintrustState; @@ -3805,6 +3864,212 @@ export type FullInitDatasetOptions = { project?: string; } & InitDatasetOptions; +async function getDatasetSnapshots( + params: + | { + state: BraintrustState; + datasetId: string; + } + | ({ + state: BraintrustState; + datasetId: string; + } & DatasetSnapshotLookup), +): Promise { + const { state, datasetId } = params; + return datasetSnapshotSchema.array().parse( + await state.appConn().post_json("api/dataset_snapshot/get", { + dataset_id: datasetId, + ...("snapshotName" in params ? { name: params.snapshotName } : {}), + ...("xactId" in params ? { xact_id: params.xactId } : {}), + }), + ); +} + +async function getDatasetSnapshot( + params: { + state: BraintrustState; + datasetId: string; + } & DatasetSnapshotLookup, +): Promise { + assertDatasetSnapshotLookup(params); + const snapshots = await getDatasetSnapshots(params); + if (snapshots.length > 1) { + throw new Error( + isDatasetSnapshotNameLookup(params) + ? `Expected a unique dataset snapshot named "${params.snapshotName}" for ${params.datasetId}` + : `Expected a unique dataset snapshot for xact_id "${params.xactId}" in ${params.datasetId}`, + ); + } + return snapshots[0]; +} + +function normalizeDatasetSelection({ + version, + environment, + snapshotName, +}: DatasetSelection): DatasetSelection { + if (version !== undefined) { + return { version }; + } + + if (snapshotName !== undefined) { + return { snapshotName }; + } + + if (environment !== undefined) { + return { environment }; + } + + return {}; +} + +async function resolveDatasetSnapshotName({ + state, + datasetId, + snapshotName, +}: { + state: BraintrustState; + datasetId: string; + snapshotName: string; +}): Promise { + const match = await getDatasetSnapshot({ + state, + datasetId, + snapshotName, + }); + if (match === undefined) { + throw new Error( + `Dataset snapshot "${snapshotName}" not found for ${datasetId}`, + ); + } + return match.xact_id; +} + +async function resolveDatasetSnapshotNameForMetadata({ + state, + lazyMetadata, + snapshotName, +}: { + state: BraintrustState; + lazyMetadata: LazyValue; + snapshotName: string; +}): Promise { + const metadata = await lazyMetadata.get(); + return await resolveDatasetSnapshotName({ + state, + datasetId: metadata.dataset.id, + snapshotName, + }); +} + +async function resolveDatasetEnvironment({ + state, + datasetId, + environment, +}: { + state: BraintrustState; + datasetId: string; + environment: string; +}): Promise { + const environmentObjectPath = `environment-object/dataset/${datasetId}/${encodeURIComponent(environment)}`; + const response = + state.orgName == null + ? await state.apiConn().get_json(environmentObjectPath) + : await state.apiConn().get_json(environmentObjectPath, { + org_name: state.orgName, + }); + return z.object({ object_version: z.string() }).parse(response) + .object_version; +} + +async function resolveDatasetEnvironmentForMetadata({ + state, + lazyMetadata, + environment, +}: { + state: BraintrustState; + lazyMetadata: LazyValue; + environment: string; +}): Promise { + const metadata = await lazyMetadata.get(); + return await resolveDatasetEnvironment({ + state, + datasetId: metadata.dataset.id, + environment, + }); +} + +async function serializeDatasetForExperiment({ + dataset, + state, +}: { + dataset: AnyDataset | DatasetRef; + state: BraintrustState; +}): Promise<{ datasetId: string; datasetVersion?: string }> { + if (!Dataset.isDataset(dataset)) { + const selection = normalizeDatasetSelection(dataset); + + if (selection.version !== undefined) { + return { + datasetId: dataset.id, + datasetVersion: selection.version, + }; + } + + if (selection.snapshotName !== undefined) { + return { + datasetId: dataset.id, + datasetVersion: await resolveDatasetSnapshotName({ + state, + datasetId: dataset.id, + snapshotName: selection.snapshotName, + }), + }; + } + + if (selection.environment !== undefined) { + return { + datasetId: dataset.id, + datasetVersion: await resolveDatasetEnvironment({ + state, + datasetId: dataset.id, + environment: selection.environment, + }), + }; + } + + return { + datasetId: dataset.id, + }; + } + + const evalData = await dataset.toEvalData(); + const selection = normalizeDatasetSelection({ + version: evalData.dataset_version, + environment: evalData.dataset_environment, + snapshotName: evalData.dataset_snapshot_name, + }); + + if (selection.version !== undefined) { + return { + datasetId: evalData.dataset_id, + datasetVersion: selection.version, + }; + } + + const datasetVersion = await dataset.version(); + if (datasetVersion !== undefined) { + return { + datasetId: evalData.dataset_id, + datasetVersion, + }; + } + + return { + datasetId: evalData.dataset_id, + }; +} + /** * Create a new dataset in a specified project. If the project does not exist, it will be created. * @@ -3812,6 +4077,9 @@ export type FullInitDatasetOptions = { * @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.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. 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. @@ -3868,6 +4136,8 @@ export function initDataset< dataset, description, version, + snapshotName, + environment, appUrl, apiKey, orgName, @@ -3879,6 +4149,14 @@ export function initDataset< state: stateArg, _internal_btql, } = options; + const selection = normalizeDatasetSelection({ + version, + environment, + snapshotName, + }); + const normalizedVersion = selection.version; + const normalizedEnvironment = selection.environment; + const normalizedSnapshotName = selection.snapshotName; const state = stateArg ?? _globalState; @@ -3919,13 +4197,57 @@ export function initDataset< }, ); - return new Dataset( + const resolvedVersion = + normalizedVersion !== undefined + ? normalizedVersion + : normalizedSnapshotName !== undefined + ? new LazyValue(async () => { + return await resolveDatasetSnapshotNameForMetadata({ + state, + lazyMetadata, + snapshotName: normalizedSnapshotName, + }); + }) + : normalizedEnvironment !== undefined + ? new LazyValue(async () => { + return await resolveDatasetEnvironmentForMetadata({ + state, + lazyMetadata, + environment: normalizedEnvironment, + }); + }) + : undefined; + + const datasetObject = new Dataset( stateArg ?? _globalState, lazyMetadata, - version, + typeof resolvedVersion === "string" ? resolvedVersion : undefined, legacy, _internal_btql, + resolvedVersion instanceof LazyValue || + normalizedEnvironment !== undefined || + normalizedSnapshotName !== undefined + ? { + ...(resolvedVersion instanceof LazyValue + ? { + lazyPinnedVersion: resolvedVersion, + } + : {}), + ...(normalizedEnvironment !== undefined + ? { + pinnedEnvironment: normalizedEnvironment, + } + : {}), + ...(normalizedSnapshotName !== undefined + ? { + pinnedSnapshotName: normalizedSnapshotName, + } + : {}), + } + : undefined, ); + + return datasetObject; } /** @@ -5706,6 +6028,18 @@ export class ObjectFetcher implements AsyncIterable< throw new Error("ObjectFetcher subclasses must have a 'getState' method"); } + protected getPinnedVersion(): string | undefined { + return this.pinnedVersion; + } + + protected setPinnedVersion(pinnedVersion: string | undefined): void { + this.pinnedVersion = pinnedVersion; + } + + protected getInternalBtql(): Record | undefined { + return this._internal_btql; + } + private async *fetchRecordsFromApi( batchSize: number | undefined, ): AsyncGenerator> { @@ -6842,6 +7176,9 @@ export class Dataset< private readonly lazyMetadata: LazyValue; private readonly __braintrust_dataset_marker = true; private newRecords = 0; + private lazyPinnedVersion?: LazyValue; + private pinnedEnvironment?: string; + private pinnedSnapshotName?: string; constructor( private state: BraintrustState, @@ -6849,6 +7186,7 @@ export class Dataset< pinnedVersion?: string, legacy?: IsLegacyDataset, _internal_btql?: Record, + pinState?: DatasetPinState, ) { // eslint-disable-next-line @typescript-eslint/consistent-type-assertions const isLegacyDataset = (legacy ?? @@ -6872,6 +7210,9 @@ export class Dataset< _internal_btql, ); this.lazyMetadata = lazyMetadata; + this.lazyPinnedVersion = pinState?.lazyPinnedVersion; + this.pinnedEnvironment = pinState?.pinnedEnvironment; + this.pinnedSnapshotName = pinState?.pinnedSnapshotName; } public get id(): Promise { @@ -6896,12 +7237,63 @@ export class Dataset< return this.state; } + public async toEvalData(): Promise<{ + dataset_id: string; + dataset_version?: string; + dataset_environment?: string; + dataset_snapshot_name?: string; + _internal_btql?: Record; + }> { + await this.getState(); + const metadata = await this.lazyMetadata.get(); + const pinnedVersion = this.getPinnedVersion(); + const internalBtql = this.getInternalBtql(); + + return { + dataset_id: metadata.dataset.id, + ...(this.pinnedEnvironment !== undefined + ? { + dataset_environment: this.pinnedEnvironment, + } + : {}), + ...(this.pinnedEnvironment === undefined && + this.pinnedSnapshotName !== undefined + ? { + dataset_snapshot_name: this.pinnedSnapshotName, + } + : {}), + ...(this.pinnedEnvironment === undefined && + this.pinnedSnapshotName === undefined && + pinnedVersion !== undefined + ? { + dataset_version: pinnedVersion, + } + : {}), + ...(internalBtql !== undefined ? { _internal_btql: internalBtql } : {}), + }; + } + protected async getState(): Promise { // Ensure the login state is populated by awaiting lazyMetadata. await this.lazyMetadata.get(); + if ( + this.lazyPinnedVersion !== undefined && + this.getPinnedVersion() === undefined + ) { + this.setPinnedVersion(await this.lazyPinnedVersion.get()); + } return this.state; } + public override async version(options?: { batchSize?: number }) { + const pinnedVersion = this.getPinnedVersion(); + if (pinnedVersion !== undefined) { + return pinnedVersion; + } + await this.getState(); + return await super.version(options); + } + private validateEvent({ metadata, expected, @@ -7078,6 +7470,116 @@ export class Dataset< return id; } + public async createSnapshot({ + name, + description, + update, + }: { + readonly name: string; + readonly description?: string; + readonly update?: boolean; + }): Promise { + await this.flush(); + const state = await this.getState(); + const datasetId = await this.id; + const currentVersion = await this.version(); + if (currentVersion === undefined) { + throw new Error("Cannot create snapshot: dataset has no version"); + } + const response = await state + .appConn() + .post_json("api/dataset_snapshot/register", { + dataset_id: datasetId, + dataset_snapshot_name: name, + description, + xact_id: currentVersion, + update, + }); + return datasetSnapshotRegisterResponseSchema.parse(response) + .dataset_snapshot; + } + + public async listSnapshots(): Promise { + const state = await this.getState(); + return await getDatasetSnapshots({ + state, + datasetId: await this.id, + }); + } + + public async getSnapshot( + lookup: DatasetSnapshotLookup, + ): Promise { + const state = await this.getState(); + const datasetId = await this.id; + return await getDatasetSnapshot({ + state, + datasetId, + ...lookup, + }); + } + + public async updateSnapshot( + snapshotId: string, + { + name, + description, + }: { + readonly name?: string; + readonly description?: string | null; + }, + ): Promise { + const state = await this.getState(); + return datasetSnapshotSchema.parse( + await state.appConn().post_json("api/dataset_snapshot/patch_id", { + id: snapshotId, + name, + description, + }), + ); + } + + public async deleteSnapshot(snapshotId: string): Promise { + const state = await this.getState(); + return datasetSnapshotSchema.parse( + await state.appConn().post_json("api/dataset_snapshot/delete_id", { + id: snapshotId, + }), + ); + } + + public async restorePreview({ + version, + }: { + readonly version: string; + }): Promise { + await this.flush(); + const state = await this.getState(); + const datasetId = await this.id; + return datasetRestorePreviewResultSchema.parse( + await state + .apiConn() + .post_json(`v1/dataset/${datasetId}/restore/preview`, { + version, + }), + ); + } + + public async restore({ + version, + }: { + readonly version: string; + }): Promise { + await this.flush(); + const state = await this.getState(); + const datasetId = await this.id; + return datasetRestoreResultSchema.parse( + await state.apiConn().post_json(`v1/dataset/${datasetId}/restore`, { + version, + }), + ); + } + /** * Summarize the dataset, including high level metrics about its size and other metadata. * @param summarizeData Whether to summarize the data. If false, only the metadata will be returned. diff --git a/js/tests/api-compatibility/api-compatibility.test.ts b/js/tests/api-compatibility/api-compatibility.test.ts index 80d15c45f..2fa625c1c 100644 --- a/js/tests/api-compatibility/api-compatibility.test.ts +++ b/js/tests/api-compatibility/api-compatibility.test.ts @@ -1492,8 +1492,8 @@ function areZodSchemaSignaturesCompatible( } /** - * Compares class signatures by extracting and comparing individual methods - * This handles cases where methods gain optional parameters or optional fields + * Compares class signatures using a cheap heuristic first and a parsed + * TypeScript AST fallback when the heuristic rejects the change. */ function areClassSignaturesCompatible( oldClass: string, @@ -1533,21 +1533,8 @@ function areClassSignaturesCompatible( // Also check normalized versions (removing optional markers and defaults) // If normalized versions are similar, the changes are likely just optionality - const normalizeClass = (classText: string): string => { - let normalized = classText; - // Remove optional markers: field?: Type -> field: Type - normalized = normalized.replace( - /([a-zA-Z_$][a-zA-Z0-9_$]*)\?:\s*/g, - "$1: ", - ); - // Remove default values - normalized = normalized.replace(/=\s*\{[^}]*\}/g, "= {}"); - normalized = normalized.replace(/=\s*[^,)}]+/g, ""); - return normalized; - }; - - const oldNormalized = normalizeClass(oldClass); - const newNormalized = normalizeClass(newClass); + const oldNormalized = normalizeClassForFastPath(oldClass); + const newNormalized = normalizeClassForFastPath(newClass); // Check if normalized versions are similar (one contains significant portion of the other) const similarityThreshold = Math.min(500, oldNormalized.length * 0.5); @@ -1555,6 +1542,432 @@ function areClassSignaturesCompatible( return true; } + const oldParsed = parseClassSignature(oldClass); + const newParsed = parseClassSignature(newClass); + if (!oldParsed || !newParsed) { + return false; + } + + return areParsedClassSignaturesCompatible(oldParsed, newParsed); +} + +interface ParsedClassSignature { + readonly name: string; + readonly typeParameters: string; + readonly heritage: string; + readonly isAbstract: boolean; + readonly members: Map; +} + +interface ParsedClassParameter { + readonly type: string; + readonly optional: boolean; + readonly rest: boolean; +} + +interface ParsedClassCallableMember { + readonly kind: "callable"; + readonly key: string; + readonly typeParameters: string; + readonly params: ParsedClassParameter[]; + readonly returnType: string; +} + +interface ParsedClassPropertyMember { + readonly kind: "property"; + readonly key: string; + readonly type: string; + readonly optional: boolean; + readonly isReadonly: boolean; +} + +type ParsedClassMember = ParsedClassCallableMember | ParsedClassPropertyMember; + +function normalizeClassForFastPath(classText: string): string { + let normalized = classText; + normalized = normalized.replace(/([a-zA-Z_$][a-zA-Z0-9_$]*)\?:\s*/g, "$1: "); + normalized = normalized.replace(/=\s*\{[^}]*\}/g, "= {}"); + normalized = normalized.replace(/=\s*[^,)}]+/g, ""); + return normalized; +} + +function normalizeClassFragment(fragment: string): string { + return normalizeTypeReference(fragment.replace(/\s+/g, " ").trim()); +} + +function hasModifier(node: ts.Node, kind: ts.SyntaxKind): boolean { + if (!ts.canHaveModifiers(node)) { + return false; + } + + return !!ts.getModifiers(node)?.some((modifier) => modifier.kind === kind); +} + +function getClassMemberVisibility( + node: ts.Node, +): "public" | "protected" | "private" { + if (hasModifier(node, ts.SyntaxKind.PrivateKeyword)) { + return "private"; + } + + if (hasModifier(node, ts.SyntaxKind.ProtectedKeyword)) { + return "protected"; + } + + return "public"; +} + +function isPrivateClassMember(member: ts.ClassElement): boolean { + if (ts.isConstructorDeclaration(member)) { + return false; + } + + if ( + (ts.isPropertyDeclaration(member) || + ts.isMethodDeclaration(member) || + ts.isGetAccessorDeclaration(member) || + ts.isSetAccessorDeclaration(member)) && + member.name && + ts.isPrivateIdentifier(member.name) + ) { + return true; + } + + return getClassMemberVisibility(member) === "private"; +} + +function isParameterProperty(parameter: ts.ParameterDeclaration): boolean { + if (!ts.canHaveModifiers(parameter)) { + return false; + } + + return !!ts.getModifiers(parameter)?.some((modifier) => { + return ( + modifier.kind === ts.SyntaxKind.PublicKeyword || + modifier.kind === ts.SyntaxKind.ProtectedKeyword || + modifier.kind === ts.SyntaxKind.PrivateKeyword || + modifier.kind === ts.SyntaxKind.ReadonlyKeyword + ); + }); +} + +function addParsedClassMember( + members: Map, + member: ParsedClassMember, +): void { + const existing = members.get(member.key); + if (existing) { + existing.push(member); + return; + } + + members.set(member.key, [member]); +} + +function parseClassParameters( + parameters: readonly ts.ParameterDeclaration[], + sourceFile: ts.SourceFile, +): ParsedClassParameter[] { + return parameters.map((parameter) => ({ + type: normalizeClassFragment(parameter.type?.getText(sourceFile) ?? ""), + optional: !!parameter.questionToken || parameter.initializer !== undefined, + rest: parameter.dotDotDotToken !== undefined, + })); +} + +function parseClassCallableMember( + key: string, + declaration: ts.SignatureDeclarationBase, + sourceFile: ts.SourceFile, + returnType: string, +): ParsedClassCallableMember { + return { + kind: "callable", + key, + typeParameters: normalizeClassFragment( + declaration.typeParameters + ?.map((typeParameter) => typeParameter.getText(sourceFile)) + .join(", ") ?? "", + ), + params: parseClassParameters(declaration.parameters, sourceFile), + returnType: normalizeClassFragment(returnType), + }; +} + +function parseClassPropertyMember( + key: string, + type: string, + optional: boolean, + isReadonly: boolean, +): ParsedClassPropertyMember { + return { + kind: "property", + key, + type: normalizeClassFragment(type), + optional, + isReadonly, + }; +} + +function parseConstructorParameterProperties( + constructor: ts.ConstructorDeclaration, + sourceFile: ts.SourceFile, +): ParsedClassPropertyMember[] { + const properties: ParsedClassPropertyMember[] = []; + + for (const parameter of constructor.parameters) { + if (!isParameterProperty(parameter)) { + continue; + } + + const visibility = getClassMemberVisibility(parameter); + if (visibility === "private") { + continue; + } + + properties.push( + parseClassPropertyMember( + `property:false:${visibility}:false:${parameter.name.getText(sourceFile)}`, + parameter.type?.getText(sourceFile) ?? "", + !!parameter.questionToken || parameter.initializer !== undefined, + hasModifier(parameter, ts.SyntaxKind.ReadonlyKeyword), + ), + ); + } + + return properties; +} + +function parseClassSignature(classText: string): ParsedClassSignature | null { + const sourceFile = ts.createSourceFile( + "class.ts", + classText, + ts.ScriptTarget.Latest, + true, + ts.ScriptKind.TS, + ); + + let declaration: ts.ClassDeclaration | undefined; + ts.forEachChild(sourceFile, (node) => { + if (ts.isClassDeclaration(node) && !declaration) { + declaration = node; + } + }); + + if (!declaration || !declaration.name) { + return null; + } + + const members = new Map(); + + for (const member of declaration.members) { + if (ts.isConstructorDeclaration(member)) { + addParsedClassMember( + members, + parseClassCallableMember( + `constructor:${getClassMemberVisibility(member)}`, + member, + sourceFile, + "", + ), + ); + for (const property of parseConstructorParameterProperties( + member, + sourceFile, + )) { + addParsedClassMember(members, property); + } + continue; + } + + if (isPrivateClassMember(member)) { + continue; + } + + const isStatic = hasModifier(member, ts.SyntaxKind.StaticKeyword); + const visibility = getClassMemberVisibility(member); + const isAbstract = hasModifier(member, ts.SyntaxKind.AbstractKeyword); + + if (ts.isPropertyDeclaration(member) && member.name) { + addParsedClassMember( + members, + parseClassPropertyMember( + `property:${isStatic}:${visibility}:${isAbstract}:${member.name.getText(sourceFile)}`, + member.type?.getText(sourceFile) ?? "", + !!member.questionToken, + hasModifier(member, ts.SyntaxKind.ReadonlyKeyword), + ), + ); + continue; + } + + if (ts.isMethodDeclaration(member)) { + addParsedClassMember( + members, + parseClassCallableMember( + `method:${isStatic}:${visibility}:${isAbstract}:${member.name.getText(sourceFile)}`, + member, + sourceFile, + member.type?.getText(sourceFile) ?? "", + ), + ); + continue; + } + + if (ts.isGetAccessorDeclaration(member)) { + addParsedClassMember( + members, + parseClassCallableMember( + `get:${isStatic}:${visibility}:${isAbstract}:${member.name.getText(sourceFile)}`, + member, + sourceFile, + member.type?.getText(sourceFile) ?? "", + ), + ); + continue; + } + + if (ts.isSetAccessorDeclaration(member)) { + addParsedClassMember( + members, + parseClassCallableMember( + `set:${isStatic}:${visibility}:${isAbstract}:${member.name.getText(sourceFile)}`, + member, + sourceFile, + "", + ), + ); + continue; + } + + if (ts.isIndexSignatureDeclaration(member)) { + addParsedClassMember( + members, + parseClassCallableMember( + `index:${visibility}:${isAbstract}`, + member, + sourceFile, + member.type?.getText(sourceFile) ?? "", + ), + ); + } + } + + return { + name: declaration.name.text, + typeParameters: normalizeClassFragment( + declaration.typeParameters + ?.map((typeParameter) => typeParameter.getText(sourceFile)) + .join(", ") ?? "", + ), + heritage: normalizeClassFragment( + declaration.heritageClauses + ?.map((heritageClause) => heritageClause.getText(sourceFile)) + .join(" ") ?? "", + ), + isAbstract: hasModifier(declaration, ts.SyntaxKind.AbstractKeyword), + members, + }; +} + +function areParsedClassSignaturesCompatible( + oldParsed: ParsedClassSignature, + newParsed: ParsedClassSignature, +): boolean { + if ( + oldParsed.name !== newParsed.name || + oldParsed.typeParameters !== newParsed.typeParameters || + oldParsed.heritage !== newParsed.heritage || + oldParsed.isAbstract !== newParsed.isAbstract + ) { + return false; + } + + for (const [key, oldMembers] of oldParsed.members) { + const newMembers = newParsed.members.get(key); + if (!newMembers) { + return false; + } + + const matchedNewMembers = new Set(); + for (const oldMember of oldMembers) { + const matchIndex = newMembers.findIndex((newMember, index) => { + return ( + !matchedNewMembers.has(index) && + areParsedClassMembersCompatible(oldMember, newMember) + ); + }); + + if (matchIndex === -1) { + return false; + } + + matchedNewMembers.add(matchIndex); + } + } + + return true; +} + +function areParsedClassMembersCompatible( + oldMember: ParsedClassMember, + newMember: ParsedClassMember, +): boolean { + if (oldMember.kind !== newMember.kind) { + return false; + } + + if (oldMember.kind === "property" && newMember.kind === "property") { + if (oldMember.isReadonly !== newMember.isReadonly) { + return false; + } + + if (oldMember.optional !== newMember.optional) { + return false; + } + + if (oldMember.type === newMember.type) { + return true; + } + + return isUnionTypeWidening(oldMember.type, newMember.type); + } + + if (oldMember.kind === "callable" && newMember.kind === "callable") { + if ( + oldMember.typeParameters !== newMember.typeParameters || + oldMember.returnType !== newMember.returnType + ) { + return false; + } + + for (let i = 0; i < oldMember.params.length; i++) { + const oldParam = oldMember.params[i]; + const newParam = newMember.params[i]; + + if (!newParam) { + return false; + } + + if (oldParam.type !== newParam.type || oldParam.rest !== newParam.rest) { + return false; + } + + if (oldParam.optional && !newParam.optional) { + return false; + } + } + + for (let i = oldMember.params.length; i < newMember.params.length; i++) { + const newParam = newMember.params[i]; + if (!newParam.optional && !newParam.rest) { + return false; + } + } + + return true; + } + return false; } @@ -1832,6 +2245,93 @@ describe("areInterfaceSignaturesCompatible", () => { }); }); +describe("areClassSignaturesCompatible", () => { + test("should allow adding private members and an optional constructor parameter", () => { + const oldClass = `export declare class Dataset { + private readonly lazyMetadata; + private readonly __braintrust_dataset_marker; + constructor( + private state: BraintrustState, + lazyMetadata: LazyValue, + pinnedVersion?: string, + ); + version(options?: { batchSize?: number }): Promise; + }`; + const newClass = `export declare class Dataset { + private readonly lazyMetadata; + private lazyPinnedVersion?: LazyValue; + private pinnedEnvironment?: string; + private pinnedSnapshotName?: string; + private readonly __braintrust_dataset_marker; + constructor( + private state: BraintrustState, + lazyMetadata: LazyValue, + pinnedVersion?: string, + pinState?: DatasetPinState, + ); + version(options?: { batchSize?: number }): Promise; + }`; + + expect(areClassSignaturesCompatible(oldClass, newClass)).toBe(true); + }); + + test("should allow adding an optional method parameter", () => { + const oldClass = `export declare class Example { + run(input: string): Promise; + }`; + const newClass = `export declare class Example { + run(input: string, options?: { retries?: number }): Promise; + }`; + + expect(areClassSignaturesCompatible(oldClass, newClass)).toBe(true); + }); + + test("should use parsed fallback when private members shift public member positions", () => { + const oldClass = `export declare class Example { + alpha(first: string, second: number): Promise<{ ok: true }>; + beta(value?: boolean): void; + }`; + const newClass = `export declare class Example { + private cacheOne?: string; + private cacheTwo?: string; + private cacheThree?: string; + alpha(first: string, second: number): Promise<{ ok: true }>; + beta(value?: boolean): void; + }`; + + expect(areClassSignaturesCompatible(oldClass, newClass)).toBe(true); + }); + + test("should ignore class member ordering", () => { + const oldClass = `export declare class Example { + alpha(): void; + protected beta(value: string): number; + }`; + const newClass = `export declare class Example { + protected beta(value: string): number; + alpha(): void; + }`; + + expect(areClassSignaturesCompatible(oldClass, newClass)).toBe(true); + }); + + test("should reject adding a required constructor parameter when fallback runs", () => { + const oldClass = `export declare class Example { + constructor(value: string); + run(input: string): void; + }`; + const newClass = `export declare class Example { + private cacheOne?: string; + private cacheTwo?: string; + private cacheThree?: string; + constructor(value: string, version: number); + run(input: string): void; + }`; + + expect(areClassSignaturesCompatible(oldClass, newClass)).toBe(false); + }); +}); + describe("areFunctionSignaturesCompatible", () => { test("should allow adding optional parameter at end", () => { const oldFn = `export function foo(a: string): void`;