fix(llmobs): use mutation response dataset versions

Author: mehulsonowal

packages/dd-trace/src/llmobs/experiments/dataset.js

@@ -449,19 +449,9 @@ class Dataset {
 
   #updateVersionFromMutationResult (result) {
     const pushedVersion = versionFromMutationResult(result)
-    if (pushedVersion !== null) {
-      this.#version = pushedVersion
-      this.#latestVersion = pushedVersion
-      return
-    }
-
-    const latestVersion = Number(this.#latestVersion)
-    if (Number.isFinite(latestVersion)) {
-      this.#version = latestVersion + 1
-      this.#latestVersion = this.#version
-    } else {
-      this.#version = null
-    }
+    if (pushedVersion === null) return
+    this.#version = pushedVersion
+    this.#latestVersion = pushedVersion
   }
 
   #addRecord (record) {

packages/dd-trace/test/llmobs/experiments/experiment.spec.js

@@ -116,7 +116,7 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     )
   })
 
-  it('advances appended dataset versions from the current latest version', async () => {
+  it('uses the dataset version from the mutation response', async () => {
     const { client: c } = clientWithMockBackend()
     const dataset = Dataset.fromExisting(c, 'demo', '', 'ds', 'proj', [], 2, 5).addRecord('a')
 
@@ -126,6 +126,27 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     assert.equal(dataset.latestVersion(), 3)
   })
 
+  it('uses each mutation response version across consecutive pushes', async () => {
+    const { client: c } = clientWithMockBackend()
+    const versions = [7, 11, 19]
+    c.batchUpdateDatasetRecords = async () => ({ records: [], version: versions.shift() })
+    const dataset = Dataset.fromExisting(c, 'demo', '', 'ds', 'proj', [], 2, 5).addRecord('before')
+
+    await dataset.push()
+    assert.equal(dataset.version(), 7)
+    assert.equal(dataset.latestVersion(), 7)
+
+    dataset.update(0, { input: 'after' })
+    await dataset.push()
+    assert.equal(dataset.version(), 11)
+    assert.equal(dataset.latestVersion(), 11)
+
+    dataset.delete(0)
+    await dataset.push()
+    assert.equal(dataset.version(), 19)
+    assert.equal(dataset.latestVersion(), 19)
+  })
+
   it('updates records added after a push and preserves concurrent edits', async () => {
     const { client: c, requests } = clientWithMockBackend()
     const dataset = new Dataset(c, 'demo').addRecord('before')
@@ -221,15 +242,15 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     assert.deepEqual(requests.at(-1).attributes.delete_records, [insertRequest.attributes.insert_records[0].id])
   })
 
-  it('handles mutation responses without a version', async () => {
+  it('keeps the current version when a mutation response omits a version', async () => {
     const { client: c } = clientWithMockBackend()
     c.batchUpdateDatasetRecords = async () => ({ records: [] })
     const dataset = Dataset.fromExisting(c, 'demo', '', 'ds', 'proj', [], null, 5).addRecord('input')
 
     await dataset.push()
 
-    assert.equal(dataset.version(), 6)
-    assert.equal(dataset.latestVersion(), 6)
+    assert.equal(dataset.version(), null)
+    assert.equal(dataset.latestVersion(), 5)
 
     const unknownVersionDataset = Dataset.fromExisting(c, 'unknown', '', 'unknown-ds', 'proj', [], null, 'unknown')
       .addRecord('input')