dd-trace 5.121.0 → 5.122.0

package/packages/dd-trace/src/llmobs/experiments/client.js

@@ -3,6 +3,9 @@
 // Control-plane HTTP client for LLM Obs Experiments. Uses the global `fetch`,
 // so this module adds no new dependency; credentials and site come from config.
 
+const { Dataset, DatasetRecord } = require('./dataset')
+const { ExperimentResult } = require('./result')
+
 const API_BASE_PATH = '/api/v2/llm-obs/v1'
 
 // Control-plane host for a Datadog site, e.g.
@@ -22,12 +25,67 @@ function appHost (site) {
   return site.split('.').length === 2 ? `app.${site}` : site
 }
 
+function datasetRecordFromResource (resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  const id = String(resource?.id ?? attrs.id ?? '')
+  if (id === '') throw new Error('Dataset record response is missing an id')
+  return new DatasetRecord(
+    attrs.input ?? null,
+    attrs.expected_output ?? null,
+    attrs.metadata ?? {},
+    id
+  )
+}
+
+function datasetVersionFromResource (resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  return attrs.valid_from_version ?? attrs.version ?? null
+}
+
+function datasetVersionFromResources (resources) {
+  const versions = resources
+    .map(datasetVersionFromResource)
+    .filter(version => version != null)
+    .map(Number)
+    .filter(Number.isFinite)
+  if (versions.length === 0) return null
+  return Math.max(...versions)
+}
+
+function datasetMutationResultFromResources (resources) {
+  return {
+    records: resources.map(datasetRecordFromResource),
+    version: datasetVersionFromResources(resources),
+  }
+}
+
+function datasetFromResource (client, projectId, resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  const version = attrs.current_version ?? null
+  return Dataset.fromExisting(
+    client,
+    String(attrs.name ?? ''),
+    String(attrs.description ?? ''),
+    resource?.id ?? attrs.id ?? null,
+    projectId,
+    [],
+    version,
+    version
+  )
+}
+
+function experimentFromResource (client, resource) {
+  const id = resource?.id
+  return new ExperimentResult(id, [], id == null ? null : `${client.appBase}/llm/experiments/${id}`)
+}
+
 class ExperimentsClient {
   #apiKey
   #appKey
   #site
   #projectName
   #timeout
+  apiBase
   #cachedProjectId
 
   constructor ({ apiKey, appKey, site, projectName, timeout = 30_000 } = {}) {
@@ -36,6 +94,7 @@ class ExperimentsClient {
     this.#site = site
     this.#projectName = projectName
     this.#timeout = timeout
+    this.apiBase = `https://${apiHost(this.#site)}`
     this.#cachedProjectId = null
   }
 
@@ -61,7 +120,7 @@ class ExperimentsClient {
   // Low-level request. Builds https://api.<site><path>, attaches both keys, and
   // returns the parsed JSON body. Throws with status + body on a non-2xx.
   async request (method, path, body) {
-    const url = `https://${apiHost(this.#site)}${path}`
+    const url = `${this.apiBase}${path}`
     const headers = {
       'DD-API-KEY': this.#apiKey,
       'DD-APPLICATION-KEY': this.#appKey,
@@ -92,6 +151,107 @@ class ExperimentsClient {
     return text ? JSON.parse(text) : {}
   }
 
+  jsonApiRequest (method, path, type, attributes) {
+    return this.request(method, path, {
+      data: { type, attributes },
+    })
+  }
+
+  async createProject (name) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/projects`, 'projects', { name })
+    return response?.data ?? null
+  }
+
+  async createDataset (projectId, attributes) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/${projectId}/datasets`, 'datasets', attributes)
+    return datasetFromResource(this, projectId, response?.data ?? null)
+  }
+
+  deleteDataset (projectId, datasetId) {
+    return this.jsonApiRequest('POST', `${API_BASE_PATH}/${projectId}/datasets/delete`, 'datasets', {
+      type: 'soft',
+      dataset_ids: [datasetId],
+    })
+  }
+
+  async listDatasets (projectId, options = {}) {
+    const query = new URLSearchParams()
+    if (options.name !== undefined) query.set('filter[name]', options.name)
+    const response = await this.request('GET', `${API_BASE_PATH}/${projectId}/datasets?${query.toString()}`)
+    const resources = Array.isArray(response?.data) ? response.data : []
+    return resources.map(resource => datasetFromResource(this, projectId, resource))
+  }
+
+  async appendDatasetRecords (projectId, datasetId, records) {
+    const response = await this.jsonApiRequest(
+      'POST',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records`,
+      'datasets',
+      { 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 resources = Array.isArray(response?.records)
+      ? response.records
+      : (Array.isArray(response?.data) ? response.data : [])
+    return resources.map(datasetRecordFromResource)
+  }
+
+  async batchUpdateDatasetRecords (projectId, datasetId, attributes) {
+    const response = await this.request(
+      'POST',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/batch_update`,
+      {
+        data: {
+          type: 'datasets',
+          id: datasetId,
+          attributes: {
+            insert_records: attributes.insert_records ?? [],
+            update_records: attributes.update_records ?? [],
+            delete_records: attributes.delete_records ?? [],
+            deduplicate: attributes.deduplicate !== false,
+            create_new_version: attributes.create_new_version !== false,
+          },
+        },
+      }
+    )
+    const resources = Array.isArray(response?.records)
+      ? response.records
+      : (Array.isArray(response?.data) ? response.data : [])
+    return datasetMutationResultFromResources(resources)
+  }
+
+  async listDatasetRecords (projectId, datasetId, options = {}) {
+    const query = new URLSearchParams()
+    if (options.cursor) query.set('page[cursor]', options.cursor)
+    if (options.version !== undefined && options.version !== null) query.set('filter[version]', String(options.version))
+    const response = await this.request(
+      'GET',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records?${query.toString()}`
+    )
+    const records = Array.isArray(response?.data) ? response.data.map(datasetRecordFromResource) : []
+    return { records, after: response?.meta?.after ?? '' }
+  }
+
+  async createExperiment (attributes) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/experiments`, 'experiments', attributes)
+    return experimentFromResource(this, response?.data ?? null)
+  }
+
+  postExperimentEvents (experimentId, attributes) {
+    return this.jsonApiRequest(
+      'POST',
+      `${API_BASE_PATH}/experiments/${experimentId}/events`,
+      'experiments',
+      attributes
+    )
+  }
+
+  updateExperiment (experimentId, attributes) {
+    return this.jsonApiRequest('PATCH', `${API_BASE_PATH}/experiments/${experimentId}`, 'experiments', attributes)
+  }
+
   // Resolve the project id for `name`, creating it if absent. The create
   // endpoint is get-or-create on name, so repeated calls return the same id.
   // Cached after the first resolution.
@@ -100,14 +260,12 @@ class ExperimentsClient {
 
     let response
     try {
-      response = await this.request('POST', `${API_BASE_PATH}/projects`, {
-        data: { type: 'projects', attributes: { name } },
-      })
+      response = await this.createProject(name)
     } catch (err) {
       throw new Error(`Failed to create or get project '${name}': ${err.message}`)
     }
 
-    this.#cachedProjectId = response?.data?.id ?? null
+    this.#cachedProjectId = response?.id ?? null
     return this.#cachedProjectId
   }
 }

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

@@ -1,73 +1,98 @@
 'use strict'
 
-const { API_BASE_PATH } = require('./client')
+const { randomUUID } = require('node:crypto')
 
-// Dataset record: { input, expectedOutput?, metadata?, id? }.
-// `id` may be user-provided before push or filled from the backend-created record.
+// Dataset record: { input, expectedOutput?, metadata?, id }.
+// IDs are generated locally unless the caller supplies one.
 class DatasetRecord {
   constructor (input, expectedOutput = null, metadata = {}, id = null) {
+    if (id != null && (typeof id !== 'string' || id.length === 0)) {
+      throw new Error('record id must be a non-empty string')
+    }
     this.input = input
     this.expectedOutput = expectedOutput ?? null
     this.metadata = metadata ?? {}
-    this.id = id ?? null
+    this.id = id ?? randomUUID()
   }
 }
 
-function createdRecordsFromResponse (response) {
-  if (Array.isArray(response?.records)) return response.records
-  if (Array.isArray(response?.data)) return response.data
-  return []
+function versionFromMutationResult (result) {
+  return result?.version ?? null
 }
 
-function recordIdFromCreatedRecord (record) {
-  return String(record?.id ?? record?.attributes?.id ?? '')
+function serializedRecord (record) {
+  return {
+    id: record.id,
+    input: record.input,
+    expected_output: record.expectedOutput ?? null,
+    metadata: record.metadata ?? {},
+  }
+}
+
+function serializedRecordUpdate (update) {
+  const output = { id: update.id }
+  if (Object.hasOwn(update, 'input') && update.input !== undefined) output.input = update.input
+  if (Object.hasOwn(update, 'expectedOutput') && update.expectedOutput !== undefined) {
+    output.expected_output = update.expectedOutput
+  }
+  if (Object.hasOwn(update, 'metadata') && update.metadata !== undefined) output.metadata = update.metadata
+  return output
 }
 
-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)
+function valuesAreEqual (left, right) {
+  return JSON.stringify(left) === JSON.stringify(right)
 }
 
-// A local buffer of dataset records, created remotely and pushed on first run
-// (or eagerly via push()). Pushes are incremental.
+function updateFromInsertedRecord (recordId, record, payload) {
+  const update = { id: recordId }
+  if (!valuesAreEqual(record.input, payload.input)) update.input = record.input
+  if (!valuesAreEqual(record.expectedOutput, payload.expected_output)) {
+    update.expectedOutput = record.expectedOutput
+  }
+  if (!valuesAreEqual(record.metadata, payload.metadata)) update.metadata = record.metadata
+  return update
+}
+
+// A local dataset model that tracks pending changes and pushes them in one batch.
 class Dataset {
   #client
   #name
   #description
   #records
-  #recordIds
+  #recordsById
+  #newRecordsById
+  #updatedRecordsById
+  #deletedRecordIds
   #id
   #projectId
-  #pushedCount
   #version
   #latestVersion
+  #pushPromise
 
   constructor (client, name, description = '') {
     this.#client = client
     this.#name = name
     this.#description = description
     this.#records = []
-    this.#recordIds = []
+    this.#recordsById = new Map()
+    this.#newRecordsById = new Map()
+    this.#updatedRecordsById = new Map()
+    this.#deletedRecordIds = new Set()
     this.#id = null
     this.#projectId = null
-    this.#pushedCount = 0
     this.#version = null
     this.#latestVersion = null
+    this.#pushPromise = Promise.resolve()
   }
 
   // Build a Dataset that already exists remotely (used by pullDataset).
-  static fromExisting (client, name, description, id, projectId, records, recordIds, version, latestVersion) {
+  static fromExisting (client, name, description, id, projectId, records, version, latestVersion) {
     const dataset = new Dataset(client, name, description)
     dataset.#id = id
     dataset.#projectId = projectId
-    dataset.#records.push(...records)
-    dataset.#recordIds.push(...recordIds)
-    dataset.#pushedCount = records.length
     dataset.#version = version ?? null
     dataset.#latestVersion = latestVersion ?? version ?? null
+    for (const record of records) dataset.#addExistingRecord(record)
     return dataset
   }
 
@@ -76,7 +101,53 @@ class Dataset {
     const record = recordOrInput instanceof DatasetRecord
       ? recordOrInput
       : new DatasetRecord(recordOrInput, expectedOutput, metadata)
-    this.#records.push(record)
+    this.#addRecord(record)
+    return this
+  }
+
+  /**
+   * Update an existing dataset record. New records are updated in place and are sent with their insert.
+   * @param {number} index Dataset record index.
+   * @param {{input?: unknown, expectedOutput?: unknown, metadata?: object}} fields Fields to update.
+   * @returns {Dataset} This dataset for chaining.
+   */
+  update (index, fields) {
+    if (fields == null || typeof fields !== 'object' || Array.isArray(fields)) {
+      throw new TypeError('record update must be an object')
+    }
+    const record = this.#recordAt(index)
+    const fieldNames = ['input', 'expectedOutput', 'metadata']
+    const providedFields = fieldNames.filter(field => Object.hasOwn(fields, field) && fields[field] !== undefined)
+    if (providedFields.length === 0) {
+      throw new Error('record update must include input, expectedOutput, or metadata')
+    }
+
+    for (const field of providedFields) {
+      record[field] = field === 'metadata' ? (fields[field] ?? {}) : fields[field]
+    }
+
+    if (this.#newRecordsById.has(record.id)) return this
+
+    const update = this.#updatedRecordsById.get(record.id) ?? { id: record.id }
+    for (const field of providedFields) update[field] = record[field]
+    this.#updatedRecordsById.set(record.id, update)
+    return this
+  }
+
+  /**
+   * Delete a dataset record. New records are removed locally without a backend operation.
+   * @param {number} index Dataset record index.
+   * @returns {Dataset} This dataset for chaining.
+   */
+  delete (index) {
+    const record = this.#recordAt(index)
+    this.#records.splice(index, 1)
+    this.#recordsById.delete(record.id)
+
+    if (this.#newRecordsById.delete(record.id)) return this
+
+    this.#updatedRecordsById.delete(record.id)
+    this.#deletedRecordIds.add(record.id)
     return this
   }
 
@@ -84,12 +155,16 @@ class Dataset {
     return this.#name
   }
 
+  description () {
+    return this.#description
+  }
+
   records () {
     return [...this.#records]
   }
 
   recordIds () {
-    return [...this.#recordIds]
+    return this.#records.map(record => record.id)
   }
 
   id () {
@@ -114,92 +189,171 @@ class Dataset {
     return `${this.#client.appBase}/llm/datasets/${this.#id}`
   }
 
-  // Eagerly create the dataset (if needed) and push any unpushed records.
-  async push () {
-    const projectId = await this.#client.ensureProjectId()
-    return this.ensureCreatedAndPushed(projectId)
+  // Eagerly create the dataset (if needed) and push all pending changes.
+  push () {
+    return this.#enqueuePush(async () => {
+      const projectId = await this.#client.ensureProjectId()
+      return this.#ensureCreatedAndPushed(projectId)
+    })
   }
 
-  // Create the remote dataset if needed, then push records added since the last
-  // push. Idempotent and incremental. Resolves to { pushedCount, totalCount } for
-  // the records attempted in this call, so callers can confirm the push landed.
-  async ensureCreatedAndPushed (projectId) {
+  // Called by Experiment.run() after it has resolved the project id.
+  ensureCreatedAndPushed (projectId) {
+    return this.#enqueuePush(() => this.#ensureCreatedAndPushed(projectId))
+  }
+
+  #enqueuePush (push) {
+    const next = this.#pushPromise.then(push, push)
+    this.#pushPromise = next.then(() => {}, () => {})
+    return next
+  }
+
+  async #ensureCreatedAndPushed (projectId) {
     if (this.#id === null) {
       let response
       try {
-        response = await this.#client.request('POST', `${API_BASE_PATH}/${projectId}/datasets`, {
-          data: { type: 'datasets', attributes: { name: this.#name, description: this.#description } },
-        })
+        response = await this.#client.createDataset(projectId, { name: this.#name, description: this.#description })
       } catch (err) {
         throw new Error(`Failed to create dataset '${this.#name}': ${err.message}`)
       }
-      this.#id = response?.data?.id ?? null
+      this.#id = response?.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
+      this.#version = response.version() ?? this.#version
+      this.#latestVersion = response.latestVersion() ?? this.#latestVersion
     }
 
-    if (this.#pushedCount >= this.#records.length) return { pushedCount: 0, totalCount: 0 }
+    const pending = this.#pendingBatch()
+    if (pending.totalCount === 0) return { pushedCount: 0, totalCount: 0 }
+
+    let result
+    try {
+      result = await this.#client.batchUpdateDatasetRecords(projectId, this.#id, pending.attributes)
+    } catch (err) {
+      throw new Error(`Failed to push changes to dataset '${this.#name}': ${err.message}`)
+    }
 
-    const pending = this.#records.slice(this.#pushedCount)
-    const records = pending.map((rec) => {
-      const out = { input: rec.input }
-      if (rec.id != null) {
-        out.id = rec.id
+    this.#updateVersionFromMutationResult(result)
+    this.#clearCommittedChanges(pending)
+    return { pushedCount: pending.totalCount, totalCount: pending.totalCount }
+  }
+
+  #pendingBatch () {
+    const insertRecords = []
+    const insertPayloads = new Map()
+    for (const [recordId, record] of this.#newRecordsById) {
+      const payload = serializedRecord(record)
+      insertRecords.push(payload)
+      insertPayloads.set(recordId, payload)
+    }
+
+    const updateRecords = []
+    const updatePayloads = new Map()
+    for (const [recordId, update] of this.#updatedRecordsById) {
+      const payload = serializedRecordUpdate(update)
+      updateRecords.push(payload)
+      updatePayloads.set(recordId, payload)
+    }
+
+    const deleteRecordIds = [...this.#deletedRecordIds]
+    return {
+      attributes: {
+        insert_records: insertRecords,
+        update_records: updateRecords,
+        delete_records: deleteRecordIds,
+        deduplicate: true,
+        create_new_version: true,
+      },
+      deleteRecordIds,
+      insertPayloads,
+      updatePayloads,
+      totalCount: new Set([
+        ...insertPayloads.keys(),
+        ...updatePayloads.keys(),
+        ...deleteRecordIds,
+      ]).size,
+    }
+  }
+
+  #clearCommittedChanges (pending) {
+    for (const [recordId, payload] of pending.insertPayloads) {
+      const current = this.#newRecordsById.get(recordId)
+      if (!current) {
+        this.#deletedRecordIds.add(recordId)
+        continue
       }
-      if (rec.expectedOutput !== null && rec.expectedOutput !== undefined) {
-        out.expected_output = rec.expectedOutput
+      if (valuesAreEqual(serializedRecord(current), payload)) {
+        this.#newRecordsById.delete(recordId)
+        continue
       }
-      if (rec.metadata && Object.keys(rec.metadata).length > 0) {
-        out.metadata = rec.metadata
+
+      this.#newRecordsById.delete(recordId)
+      const update = this.#updatedRecordsById.get(recordId) ??
+        updateFromInsertedRecord(recordId, current, payload)
+      this.#updatedRecordsById.set(recordId, update)
+    }
+
+    for (const [recordId, payload] of pending.updatePayloads) {
+      const current = this.#updatedRecordsById.get(recordId)
+      if (current && valuesAreEqual(serializedRecordUpdate(current), payload)) {
+        this.#updatedRecordsById.delete(recordId)
       }
-      return out
-    })
+    }
 
-    let response
-    try {
-      response = await this.#client.request(
-        'POST',
-        `${API_BASE_PATH}/${projectId}/datasets/${this.#id}/records`,
-        { data: { type: 'datasets', attributes: { records } } }
-      )
-    } catch (err) {
-      throw new Error(`Failed to push records to dataset '${this.#name}': ${err.message}`)
+    for (const recordId of pending.deleteRecordIds) {
+      if (!this.#recordsById.has(recordId)) this.#deletedRecordIds.delete(recordId)
     }
+  }
 
-    // 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 {
+  #updateVersionFromMutationResult (result) {
+    const pushedVersion = versionFromMutationResult(result)
+    if (pushedVersion !== null) {
       this.#version = pushedVersion
-      this.#latestVersion = Math.max(Number(this.#latestVersion ?? pushedVersion), pushedVersion)
+      this.#latestVersion = pushedVersion
+      return
     }
 
-    let pushedCount = 0
-    for (const [index, node] of created.entries()) {
-      const recordId = recordIdFromCreatedRecord(node)
-      if (recordId !== '') {
-        pushedCount++
-        pending[index].id = recordId
-      }
-      this.#recordIds.push(recordId)
+    const latestVersion = Number(this.#latestVersion)
+    if (Number.isFinite(latestVersion)) {
+      this.#version = latestVersion + 1
+      this.#latestVersion = this.#version
+    } else {
+      this.#version = null
     }
-    for (let i = created.length; i < pending.length; i++) this.#recordIds.push('')
+  }
+
+  #addRecord (record) {
+    if (this.#recordsById.has(record.id)) throw new Error(`Duplicate record id '${record.id}'`)
+    this.#records.push(record)
+    this.#recordsById.set(record.id, record)
+    this.#deletedRecordIds.delete(record.id)
+    this.#newRecordsById.set(record.id, record)
+  }
 
-    // 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.
-    this.#pushedCount += pending.length
+  #addExistingRecord (record) {
+    if (record.id === null || record.id === undefined || record.id === '') {
+      throw new Error('Dataset records pulled from the backend must have an id')
+    }
+    if (!(record instanceof DatasetRecord)) {
+      record = new DatasetRecord(
+        record.input,
+        record.expectedOutput,
+        record.metadata,
+        record.id
+      )
+    }
+    if (this.#recordsById.has(record.id)) throw new Error(`Duplicate record id '${record.id}'`)
+    this.#records.push(record)
+    this.#recordsById.set(record.id, record)
+  }
 
-    return { pushedCount, totalCount: pending.length }
+  #recordAt (index) {
+    if (!Number.isInteger(index) || index < 0 || index >= this.#records.length) {
+      throw new RangeError(`Dataset record index ${index} is out of range`)
+    }
+    return this.#records[index]
   }
 }