fix(llmobs): address dataset review feedback

Author: mehulsonowal

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

@@ -2,7 +2,8 @@
 
 const { API_BASE_PATH } = require('./client')
 
-// Immutable dataset record: { input, expectedOutput?, metadata?, id? }.
+// Dataset record: { input, expectedOutput?, metadata?, id? }.
+// `id` may be user-provided before push or filled from the backend-created record.
 class DatasetRecord {
   constructor (input, expectedOutput = null, metadata = {}, id = null) {
     this.input = input
@@ -12,6 +13,24 @@ class DatasetRecord {
   }
 }
 
+function createdRecordsFromResponse (response) {
+  if (Array.isArray(response?.records)) return response.records
+  if (Array.isArray(response?.data)) return response.data
+  return []
+}
+
+function recordIdFromCreatedRecord (record) {
+  return String(record?.id ?? record?.attributes?.id ?? '')
+}
+
+function versionFromCreatedRecords (records) {
+  const versions = records
+    .map(record => Number(record?.attributes?.valid_from_version ?? record?.attributes?.version))
+    .filter(Number.isFinite)
+  if (versions.length === 0) return null
+  return Math.max(...versions)
+}
+
 // A local buffer of dataset records, created remotely and pushed on first run
 // (or eagerly via push()). Pushes are incremental.
 class Dataset {
@@ -115,6 +134,9 @@ class Dataset {
         throw new Error(`Failed to create dataset '${this.#name}': ${err.message}`)
       }
       this.#id = response?.data?.id ?? null
+      if (this.#id === null) {
+        throw new Error(`Failed to create dataset '${this.#name}': backend response is missing dataset id`)
+      }
       this.#projectId = projectId
       this.#version = response?.data?.attributes?.current_version ?? this.#version
       this.#latestVersion = response?.data?.attributes?.current_version ?? this.#latestVersion
@@ -148,20 +170,30 @@ class Dataset {
       throw new Error(`Failed to push records to dataset '${this.#name}': ${err.message}`)
     }
 
-    // The append-records response returns created records under a top-level
-    // `records` field, not the usual `data` envelope.
-    const created = response?.records
+    // The append-records response has used both a top-level `records` array
+    // and JSON:API `data` resources. Accept either so generated/custom record
+    // ids are preserved for experiment row tagging.
+    const created = createdRecordsFromResponse(response)
+    const pushedVersion = versionFromCreatedRecords(created)
+    if (pushedVersion === null) {
+      // The dataset contents changed, but the backend did not report the new
+      // version. Avoid pinning later experiments to the pre-append create version.
+      this.#version = null
+    } else {
+      this.#version = pushedVersion
+      this.#latestVersion = Math.max(Number(this.#latestVersion ?? pushedVersion), pushedVersion)
+    }
+
     let pushedCount = 0
-    if (Array.isArray(created)) {
-      for (const node of created) {
-        const recordId = String(node?.id ?? '')
-        if (recordId !== '') pushedCount++
-        this.#recordIds.push(recordId)
+    for (const [index, node] of created.entries()) {
+      const recordId = recordIdFromCreatedRecord(node)
+      if (recordId !== '') {
+        pushedCount++
+        pending[index].id = recordId
       }
-      for (let i = created.length; i < pending.length; i++) this.#recordIds.push('')
-    } else {
-      for (let i = 0; i < pending.length; i++) this.#recordIds.push('')
+      this.#recordIds.push(recordId)
     }
+    for (let i = created.length; i < pending.length; i++) this.#recordIds.push('')
 
     // Advance by the snapshotted pending count, not the live records length,
     // so records added while this push was in flight aren't skipped by the next push.

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

@@ -69,6 +69,7 @@ function recordsFromCsv (csvPath, options) {
     delimiter: csvDelimiter,
     bom: true,
     relaxColumnCount: true,
+    skipEmptyLines: true,
   })
   const header = rows.shift()
   if (!header || header.every(column => column === '')) {
@@ -98,7 +99,7 @@ function recordsFromCsv (csvPath, options) {
     for (let i = 0; i < header.length; i++) row[header[i]] = values[i] ?? ''
     return {
       inputData: selectedColumns(row, inputDataColumns),
-      expectedOutput: selectedColumns(row, expectedOutputColumns),
+      expectedOutput: expectedOutputColumns.length === 0 ? undefined : selectedColumns(row, expectedOutputColumns),
       metadata: selectedColumns(row, metadataColumns),
       id: idColumn === undefined ? undefined : row[idColumn],
     }
@@ -189,7 +190,7 @@ class Experiments {
         for (;;) {
           const query = new URLSearchParams()
           if (cursor) query.set('page[cursor]', cursor)
-          if (version !== undefined && version !== null) query.set('filter[version]', String(version))
+          if (datasetVersion !== null) query.set('filter[version]', String(datasetVersion))
           const queryString = query.toString() ? `?${query.toString()}` : ''
           // eslint-disable-next-line no-await-in-loop
           const resp = await this.#client.request(
@@ -198,8 +199,14 @@ class Experiments {
           )
           for (const item of resp?.data ?? []) {
             const attrs = item?.attributes ?? item
-            recs.push(new DatasetRecord(attrs?.input ?? null, attrs?.expected_output ?? null, attrs?.metadata ?? {}))
-            ids.push(String(item?.id ?? ''))
+            const recordId = String(item?.id ?? attrs?.id ?? '')
+            recs.push(new DatasetRecord(
+              attrs?.input ?? null,
+              attrs?.expected_output ?? null,
+              attrs?.metadata ?? {},
+              recordId === '' ? null : recordId
+            ))
+            ids.push(recordId)
           }
           cursor = resp?.meta?.after ?? ''
           if (!cursor) break

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

@@ -5,7 +5,7 @@ const { afterEach, beforeEach, describe, it } = require('mocha')
 const sinon = require('sinon')
 
 const { ExperimentsClient } = require('../../../src/llmobs/experiments/client')
-const { Dataset } = require('../../../src/llmobs/experiments/dataset')
+const { Dataset, DatasetRecord } = require('../../../src/llmobs/experiments/dataset')
 const { Experiment } = require('../../../src/llmobs/experiments/experiment')
 
 // Routes the control-plane + events calls a run makes, recording each request.
@@ -240,7 +240,6 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
   })
 
   it('experiment create body carries project_id, dataset_id, dataset_version, config and ensure_unique', async () => {
-    const { DatasetRecord } = require('../../../src/llmobs/experiments/dataset')
     const dataset = Dataset.fromExisting(
       client,
       'demo',
@@ -267,6 +266,23 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     assert.deepEqual(create.body.data.attributes.config, { approach: 'kw' })
   })
 
+  it('uses the version returned by appending records when creating an experiment', async () => {
+    installFetch(calls, {
+      'POST /api/v2/llm-obs/v1/proj/datasets': { data: { id: 'ds', attributes: { current_version: 1 } } },
+      'POST /api/v2/llm-obs/v1/proj/datasets/ds/records': {
+        data: [{ id: 'rec-0', attributes: { valid_from_version: 2 } }],
+      },
+    })
+    const dataset = new Dataset(client, 'demo').addRecord('x')
+
+    await new Experiment(client, { name: 'exp-demo', dataset, task: (i) => i }).run()
+
+    const create = calls.find(c =>
+      c.method === 'POST' && c.path.endsWith('/experiments') && !c.path.includes('/events')
+    )
+    assert.equal(create.body.data.attributes.dataset_version, 2)
+  })
+
   it('validates required options', () => {
     const dataset = new Dataset(client, 'demo')
     assert.throws(() => new Experiment(client, { dataset, task: (i) => i }), /name/)
@@ -275,7 +291,6 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
   })
 
   it('exposes dataset getters and accepts a DatasetRecord instance', () => {
-    const { DatasetRecord } = require('../../../src/llmobs/experiments/dataset')
     const dataset = new Dataset(client, 'my-name', 'desc').addRecord(new DatasetRecord('in', 'out', { m: 1 }))
     assert.equal(dataset.name(), 'my-name')
     assert.equal(dataset.id(), null)
@@ -292,12 +307,39 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     const result = await dataset.push()
     assert.deepEqual(dataset.recordIds(), ['', ''])
     assert.deepEqual(result, { pushedCount: 0, totalCount: 2 })
+    assert.equal(dataset.version(), null)
   })
 
   it('resolves with the pushed/total record counts on a successful push', async () => {
     const dataset = new Dataset(client, 'demo').addRecord('a').addRecord('b')
     const result = await dataset.push()
     assert.deepEqual(result, { pushedCount: 2, totalCount: 2 })
+    assert.deepEqual(dataset.recordIds(), ['rec-0', 'rec-1'])
+    assert.deepEqual(dataset.records().map(record => record.id), ['rec-0', 'rec-1'])
+  })
+
+  it('submits custom record ids and keeps ids from JSON:API push responses', async () => {
+    installFetch(calls, {
+      'POST /api/v2/llm-obs/v1/proj/datasets/ds/records': {
+        data: [
+          { id: 'custom-a', attributes: { valid_from_version: 2 } },
+          { id: 'custom-b', attributes: { valid_from_version: 2 } },
+        ],
+      },
+    })
+    const dataset = new Dataset(client, 'demo')
+      .addRecord(new DatasetRecord('a', null, {}, 'custom-a'))
+      .addRecord(new DatasetRecord('b', null, {}, 'custom-b'))
+
+    const result = await dataset.push()
+
+    const recordsCall = calls.find(call => call.path === '/api/v2/llm-obs/v1/proj/datasets/ds/records')
+    assert.deepEqual(recordsCall.body.data.attributes.records.map(record => record.id), ['custom-a', 'custom-b'])
+    assert.deepEqual(dataset.recordIds(), ['custom-a', 'custom-b'])
+    assert.deepEqual(dataset.records().map(record => record.id), ['custom-a', 'custom-b'])
+    assert.equal(dataset.version(), 2)
+    assert.equal(dataset.latestVersion(), 2)
+    assert.deepEqual(result, { pushedCount: 2, totalCount: 2 })
   })
 
   it('resolves with a lower pushedCount when the backend confirms fewer records than sent', async () => {
@@ -308,6 +350,8 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     const result = await dataset.push()
     assert.deepEqual(result, { pushedCount: 1, totalCount: 2 })
     assert.deepEqual(dataset.recordIds(), ['rec-0', ''])
+    assert.deepEqual(dataset.records().map(record => record.id), ['rec-0', null])
+    assert.equal(dataset.version(), null)
   })
 
   it('resolves with zero counts when there is nothing new to push', async () => {
@@ -340,12 +384,12 @@ describe('LLMObs Experiments — dataset + experiment run', () => {
     assert.equal(experiment.url(), 'https://app.datadoghq.com/llm/experiments/exp')
   })
 
-  it('throws when the dataset has no id after push', async () => {
+  it('throws when the dataset create response has no id', async () => {
     installFetch(calls, { 'POST /api/v2/llm-obs/v1/proj/datasets': { data: {} } })
     const dataset = new Dataset(client, 'demo').addRecord('x')
     await assert.rejects(
       () => new Experiment(client, { name: 'exp-demo', dataset, task: (i) => i }).run(),
-      /has no id after push/
+      /backend response is missing dataset id/
     )
   })
 

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

@@ -115,6 +115,47 @@ describe('LLMObs Experiments facade', () => {
       }
     })
 
+    it('omits expected output and skips blank lines for input-only CSV datasets', async () => {
+      const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'dd-trace-llmobs-csv-'))
+      const csvPath = path.join(dir, 'dataset.csv')
+      fs.writeFileSync(csvPath, 'question\nhello\n\nworld\n')
+
+      const calls = []
+      global.fetch.callsFake(async (url, opts) => {
+        const u = new URL(url)
+        const body = opts.body ? JSON.parse(opts.body) : undefined
+        calls.push({ method: opts.method, path: u.pathname, body })
+        let payload = {}
+        if (u.pathname === '/api/v2/llm-obs/v1/projects') {
+          payload = { data: { id: 'proj' } }
+        } else if (u.pathname === '/api/v2/llm-obs/v1/proj/datasets') {
+          payload = { data: { id: 'ds' } }
+        } else if (u.pathname === '/api/v2/llm-obs/v1/proj/datasets/ds/records') {
+          payload = { records: [{ id: 'rec-1' }, { id: 'rec-2' }] }
+        }
+        return { ok: true, status: 200, text: sinon.stub().resolves(JSON.stringify(payload)) }
+      })
+
+      try {
+        const dataset = createExperiments(enabledConfig()).createDatasetFromCsv(csvPath, 'csv-dataset', {
+          inputDataColumns: ['question'],
+        })
+        assert.equal(dataset.records().length, 2)
+        assert.deepEqual(dataset.records().map(record => record.input), [{ question: 'hello' }, { question: 'world' }])
+        assert.deepEqual(dataset.records().map(record => record.expectedOutput), [null, null])
+
+        await dataset.push()
+
+        const recordsCall = calls.find(call => call.path === '/api/v2/llm-obs/v1/proj/datasets/ds/records')
+        assert.deepEqual(recordsCall.body.data.attributes.records, [
+          { input: { question: 'hello' } },
+          { input: { question: 'world' } },
+        ])
+      } finally {
+        fs.rmSync(dir, { recursive: true, force: true })
+      }
+    })
+
     it('validates CSV headers before creating a dataset', () => {
       const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'dd-trace-llmobs-csv-'))
       const csvPath = path.join(dir, 'dataset.csv')
@@ -207,6 +248,8 @@ describe('LLMObs Experiments facade', () => {
       assert.deepEqual(ds.records()[0].input, { q: '2+2' })
       assert.equal(ds.records()[0].expectedOutput, '4')
       assert.deepEqual(ds.records()[0].metadata, { a: 1 })
+      assert.equal(ds.records()[0].id, 'r1')
+      assert.equal(ds.records()[1].id, 'r2')
     })
 
     it('passes explicit dataset version when reading records', async () => {
@@ -231,6 +274,28 @@ describe('LLMObs Experiments facade', () => {
       assert.equal(ds.latestVersion(), 7)
     })
 
+    it('pins the current version when pulling latest records', async () => {
+      global.fetch.callsFake(async (url) => {
+        const u = new URL(url)
+        let payload
+        if (u.pathname === '/api/v2/llm-obs/v1/projects') {
+          payload = { data: { id: 'proj' } }
+        } else if (u.pathname === '/api/v2/llm-obs/v1/proj/datasets') {
+          payload = { data: [{ id: 'ds9', attributes: { name: 'wanted', description: 'd', current_version: 7 } }] }
+        } else if (u.pathname === '/api/v2/llm-obs/v1/proj/datasets/ds9/records') {
+          assert.equal(u.searchParams.get('filter[version]'), '7')
+          payload = { data: [{ id: 'r1', attributes: { input: 'i1' } }] }
+        } else {
+          payload = {}
+        }
+        return { ok: true, status: 200, text: sinon.stub().resolves(JSON.stringify(payload)) }
+      })
+
+      const ds = await createExperiments(enabledConfig()).pullDataset('wanted')
+      assert.equal(ds.version(), 7)
+      assert.equal(ds.records().length, 1)
+    })
+
     it('waits (backoff) until the expected record count is readable', async () => {
       const one = { data: [{ id: 'r1', attributes: { input: 'i1' } }] }
       const two = { data: [{ id: 'r1', attributes: { input: 'i1' } }, { id: 'r2', attributes: { input: 'i2' } }] }