dd-trace 6.11.0 → 6.12.0

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

@@ -2,16 +2,30 @@
 
 const { randomUUID } = require('node:crypto')
 
-// Dataset record: { input, expectedOutput?, metadata?, id }.
+/** @typedef {{add?: string[], remove?: string[], replace?: string[]}} TagOperations */
+/**
+ * @typedef {object} PendingBatch
+ * @property {object} attributes
+ * @property {string[]} deleteRecordIds
+ * @property {Map<string, object>} insertPayloads
+ * @property {Map<string, object>} updatePayloads
+ * @property {number} totalCount
+ * @property {Map<string, TagOperations>} [inFlightTagOperations]
+ */
+
+const { tagOperationsAreEmpty, validateTagsList } = require('./util')
+
+// Dataset record: { input, expectedOutput?, metadata?, id, tags? }.
 // IDs are generated locally unless the caller supplies one.
 class DatasetRecord {
-  constructor (input, expectedOutput = null, metadata = {}, id = null) {
+  constructor (input, expectedOutput = null, metadata = {}, id = null, tags = []) {
     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.tags = validateTagsList(tags)
     this.id = id ?? randomUUID()
   }
 }
@@ -21,12 +35,14 @@ function versionFromMutationResult (result) {
 }
 
 function serializedRecord (record) {
-  return {
+  const output = {
     id: record.id,
     input: record.input,
     expected_output: record.expectedOutput ?? null,
     metadata: record.metadata ?? {},
   }
+  if (record.tags.length > 0) output.tags = record.tags
+  return output
 }
 
 function serializedRecordUpdate (update) {
@@ -36,9 +52,57 @@ function serializedRecordUpdate (update) {
     output.expected_output = update.expectedOutput
   }
   if (Object.hasOwn(update, 'metadata') && update.metadata !== undefined) output.metadata = update.metadata
+  if (Object.hasOwn(update, 'tagOperations')) {
+    output.tag_operations = serializedTagOperations(update.tagOperations)
+  }
+  return output
+}
+
+function serializedTagOperations (operations) {
+  const output = {}
+  if (Object.hasOwn(operations, 'add')) output.add = operations.add
+  if (Object.hasOwn(operations, 'remove')) output.remove = operations.remove
+  if (Object.hasOwn(operations, 'replace')) output.set = operations.replace
   return output
 }
 
+/**
+ * @param {TagOperations} operations
+ * @returns {TagOperations}
+ */
+function copyTagOperations (operations) {
+  return Object.fromEntries(Object.entries(operations).map(([key, tags]) => [key, [...tags]]))
+}
+
+function mergeTagOperations (operations, operation, tags) {
+  if (operation === 'replace') return { replace: [...tags] }
+  if (Object.hasOwn(operations, 'replace')) {
+    const replaced = new Set(operations.replace)
+    for (const tag of tags) {
+      if (operation === 'add') replaced.add(tag)
+      else replaced.delete(tag)
+    }
+    return { replace: [...replaced].sort() }
+  }
+
+  const add = new Set(operations.add)
+  const remove = new Set(operations.remove)
+  for (const tag of tags) {
+    if (operation === 'add') {
+      if (remove.has(tag)) remove.delete(tag)
+      else add.add(tag)
+    } else if (add.has(tag)) {
+      add.delete(tag)
+    } else {
+      remove.add(tag)
+    }
+  }
+  const merged = {}
+  if (add.size > 0) merged.add = [...add].sort()
+  if (remove.size > 0) merged.remove = [...remove].sort()
+  return merged
+}
+
 function valuesAreEqual (left, right) {
   return JSON.stringify(left) === JSON.stringify(right)
 }
@@ -63,21 +127,25 @@ class Dataset {
   #newRecordsById
   #updatedRecordsById
   #deletedRecordIds
+  #pendingTagOperations
   #id
   #projectId
   #version
   #latestVersion
   #pushPromise
+  #filterTags
 
-  constructor (client, name, description = '') {
+  constructor (client, name, description = '', filterTags = []) {
     this.#client = client
     this.#name = name
     this.#description = description
+    this.#filterTags = validateTagsList(filterTags)
     this.#records = []
     this.#recordsById = new Map()
     this.#newRecordsById = new Map()
     this.#updatedRecordsById = new Map()
     this.#deletedRecordIds = new Set()
+    this.#pendingTagOperations = new Map()
     this.#id = null
     this.#projectId = null
     this.#version = null
@@ -86,8 +154,8 @@ class Dataset {
   }
 
   // Build a Dataset that already exists remotely (used by pullDataset).
-  static fromExisting (client, name, description, id, projectId, records, version, latestVersion) {
-    const dataset = new Dataset(client, name, description)
+  static fromExisting (client, name, description, id, projectId, records, version, latestVersion, filterTags) {
+    const dataset = new Dataset(client, name, description, filterTags)
     dataset.#id = id
     dataset.#projectId = projectId
     dataset.#version = version ?? null
@@ -97,14 +165,70 @@ class Dataset {
   }
 
   // Append a record. Accepts a DatasetRecord or (input, expectedOutput?, metadata?).
-  addRecord (recordOrInput, expectedOutput, metadata) {
+  addRecord (recordOrInput, expectedOutput, metadata, tags) {
     const record = recordOrInput instanceof DatasetRecord
       ? recordOrInput
-      : new DatasetRecord(recordOrInput, expectedOutput, metadata)
+      : new DatasetRecord(recordOrInput, expectedOutput, metadata, null, tags)
     this.#addRecord(record)
     return this
   }
 
+  /**
+   * Add tags to a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {Dataset} This dataset for chaining.
+   */
+  addTags (index, tags) {
+    const validated = validateTagsList(tags)
+    const record = this.#recordAt(index)
+    const next = new Set(record.tags)
+    const added = []
+    for (const tag of validated) {
+      if (next.has(tag)) continue
+      next.add(tag)
+      added.push(tag)
+    }
+    record.tags = [...next].sort()
+    this.#queueTagOperation(record.id, 'add', added)
+    return this
+  }
+
+  /**
+   * Remove tags from a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {Dataset} This dataset for chaining.
+   */
+  removeTags (index, tags) {
+    const validated = validateTagsList(tags)
+    const record = this.#recordAt(index)
+    const next = new Set(record.tags)
+    const removed = []
+    for (const tag of validated) {
+      if (!next.has(tag)) continue
+      next.delete(tag)
+      removed.push(tag)
+    }
+    record.tags = [...next].sort()
+    this.#queueTagOperation(record.id, 'remove', removed)
+    return this
+  }
+
+  /**
+   * Replace all tags on a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {Dataset} This dataset for chaining.
+   */
+  replaceTags (index, tags) {
+    const validated = validateTagsList(tags)
+    const record = this.#recordAt(index)
+    record.tags = validated
+    this.#queueTagOperation(record.id, 'replace', validated)
+    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.
@@ -130,6 +254,9 @@ class Dataset {
 
     const update = this.#updatedRecordsById.get(record.id) ?? { id: record.id }
     for (const field of providedFields) update[field] = record[field]
+    if (this.#pendingTagOperations.has(record.id)) {
+      update.tagOperations = this.#pendingTagOperations.get(record.id)
+    }
     this.#updatedRecordsById.set(record.id, update)
     return this
   }
@@ -144,13 +271,47 @@ class Dataset {
     this.#records.splice(index, 1)
     this.#recordsById.delete(record.id)
 
-    if (this.#newRecordsById.delete(record.id)) return this
+    if (this.#newRecordsById.delete(record.id)) {
+      this.#pendingTagOperations.delete(record.id)
+      return this
+    }
 
     this.#updatedRecordsById.delete(record.id)
+    this.#pendingTagOperations.delete(record.id)
     this.#deletedRecordIds.add(record.id)
     return this
   }
 
+  /**
+   * Queue a tag operation for the next dataset push.
+   * @param {string} recordId Dataset record id.
+   * @param {'add' | 'remove' | 'replace'} operation Tag operation to queue.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {void}
+   */
+  #queueTagOperation (recordId, operation, tags) {
+    if (operation !== 'replace' && tags.length === 0) return
+    const operations = mergeTagOperations(this.#pendingTagOperations.get(recordId) ?? {}, operation, tags)
+    if (tagOperationsAreEmpty(operations)) {
+      this.#pendingTagOperations.delete(recordId)
+      const update = this.#updatedRecordsById.get(recordId)
+      if (update) {
+        const hasRecordUpdate = Object.hasOwn(update, 'input') ||
+          Object.hasOwn(update, 'expectedOutput') ||
+          Object.hasOwn(update, 'metadata')
+        delete update.tagOperations
+        if (!hasRecordUpdate) this.#updatedRecordsById.delete(recordId)
+      }
+      return
+    }
+    this.#pendingTagOperations.set(recordId, operations)
+    if (!this.#newRecordsById.has(recordId)) {
+      const update = this.#updatedRecordsById.get(recordId) ?? { id: recordId }
+      update.tagOperations = operations
+      this.#updatedRecordsById.set(recordId, update)
+    }
+  }
+
   name () {
     return this.#name
   }
@@ -183,6 +344,14 @@ class Dataset {
     return this.#latestVersion
   }
 
+  /**
+   * Return the tags used to filter this dataset.
+   * @returns {string[]} Dataset record filter tags.
+   */
+  filterTags () {
+    return [...this.#filterTags]
+  }
+
   // Dashboard URL for this dataset, or null until pushed/pulled.
   url () {
     if (this.#id === null) return null
@@ -228,10 +397,13 @@ class Dataset {
     const pending = this.#pendingBatch()
     if (pending.totalCount === 0) return { pushedCount: 0, totalCount: 0 }
 
+    this.#detachCommittedTagOperations(pending)
+
     let result
     try {
       result = await this.#client.batchUpdateDatasetRecords(projectId, this.#id, pending.attributes)
     } catch (err) {
+      this.#restoreFailedTagOperations(pending)
       throw new Error(`Failed to push changes to dataset '${this.#name}': ${err.message}`)
     }
 
@@ -252,6 +424,9 @@ class Dataset {
     const updateRecords = []
     const updatePayloads = new Map()
     for (const [recordId, update] of this.#updatedRecordsById) {
+      const tagOperations = this.#pendingTagOperations.get(recordId)
+      if (tagOperations) update.tagOperations = tagOperations
+      else delete update.tagOperations
       const payload = serializedRecordUpdate(update)
       updateRecords.push(payload)
       updatePayloads.set(recordId, payload)
@@ -277,6 +452,58 @@ class Dataset {
     }
   }
 
+  /**
+   * Detach tag changes sent by this batch so edits made while the request is in flight
+   * are queued relative to the response that this batch will commit.
+   * @param {PendingBatch} pending
+   * @returns {void}
+   */
+  #detachCommittedTagOperations (pending) {
+    pending.inFlightTagOperations = new Map()
+    for (const [recordId] of pending.insertPayloads) {
+      const operations = this.#pendingTagOperations.get(recordId)
+      if (!operations) continue
+      pending.inFlightTagOperations.set(recordId, copyTagOperations(operations))
+      this.#pendingTagOperations.delete(recordId)
+    }
+    for (const [recordId, payload] of pending.updatePayloads) {
+      const operations = this.#pendingTagOperations.get(recordId)
+      if (!operations || !Object.hasOwn(payload, 'tag_operations')) continue
+
+      pending.inFlightTagOperations.set(recordId, copyTagOperations(operations))
+      this.#pendingTagOperations.delete(recordId)
+      const update = this.#updatedRecordsById.get(recordId)
+      if (update) delete update.tagOperations
+    }
+  }
+
+  /**
+   * Restore tag changes when a batch request fails, including edits made while it was in flight.
+   * @param {PendingBatch} pending
+   * @returns {void}
+   */
+  #restoreFailedTagOperations (pending) {
+    if (!pending.inFlightTagOperations) return
+    for (const [recordId, operations] of pending.inFlightTagOperations) {
+      const record = this.#recordsById.get(recordId)
+      if (!record || this.#newRecordsById.has(recordId) || this.#deletedRecordIds.has(recordId)) continue
+
+      const update = this.#updatedRecordsById.get(recordId) ?? { id: recordId }
+      const queuedOperations = this.#pendingTagOperations.get(recordId)
+      const restoredOperations = queuedOperations
+        ? { replace: [...record.tags] }
+        : copyTagOperations(operations)
+      this.#pendingTagOperations.set(recordId, restoredOperations)
+      update.tagOperations = restoredOperations
+      this.#updatedRecordsById.set(recordId, update)
+    }
+  }
+
+  /**
+   * Clear the changes represented by a completed batch while retaining concurrent local edits.
+   * @param {PendingBatch} pending
+   * @returns {void}
+   */
   #clearCommittedChanges (pending) {
     for (const [recordId, payload] of pending.insertPayloads) {
       const current = this.#newRecordsById.get(recordId)
@@ -286,19 +513,29 @@ class Dataset {
       }
       if (valuesAreEqual(serializedRecord(current), payload)) {
         this.#newRecordsById.delete(recordId)
+        this.#pendingTagOperations.delete(recordId)
         continue
       }
 
       this.#newRecordsById.delete(recordId)
       const update = this.#updatedRecordsById.get(recordId) ??
         updateFromInsertedRecord(recordId, current, payload)
+      const queuedOperations = this.#pendingTagOperations.get(recordId)
+      if (queuedOperations) update.tagOperations = queuedOperations
       this.#updatedRecordsById.set(recordId, update)
     }
 
     for (const [recordId, payload] of pending.updatePayloads) {
       const current = this.#updatedRecordsById.get(recordId)
-      if (current && valuesAreEqual(serializedRecordUpdate(current), payload)) {
+      const queuedOperations = this.#pendingTagOperations.get(recordId)
+      if (!current || queuedOperations) continue
+
+      const comparison = { ...current }
+      const committedOperations = pending.inFlightTagOperations?.get(recordId)
+      if (committedOperations) comparison.tagOperations = committedOperations
+      if (valuesAreEqual(serializedRecordUpdate(comparison), payload)) {
         this.#updatedRecordsById.delete(recordId)
+        this.#pendingTagOperations.delete(recordId)
       }
     }
 
@@ -341,7 +578,8 @@ class Dataset {
         record.input,
         record.expectedOutput,
         record.metadata,
-        record.id
+        record.id,
+        record.tags
       )
     }
     if (this.#recordsById.has(record.id)) throw new Error(`Duplicate record id '${record.id}'`)

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

@@ -13,6 +13,7 @@ const {
   normalizeEvaluators,
   mergeTags,
   normalizeJsonMetricValue,
+  recordTagsToObject,
   sleep,
   stringify,
   timestampMs,
@@ -20,7 +21,7 @@ const {
 } = require('./util')
 
 // One span per experiment row (LLM Obs experiment span wire format).
-function toSpan (row, metadata, ids, spanName, userTags) {
+function toSpan (row, metadata, ids, spanName, userTags, recordTags) {
   const meta = {
     input: row.input ?? null,
     output: row.output ?? null,
@@ -33,6 +34,20 @@ function toSpan (row, metadata, ids, spanName, userTags) {
     meta.error = { type: row.errorType ?? '', message: row.errorMessage ?? '', stack: row.errorStack ?? '' }
   }
 
+  const tags = buildTags({
+    ...userTags,
+    ...recordTagsToObject(recordTags),
+  }, {
+    experiment_id: ids.experimentId,
+    run_id: ids.runId,
+    run_iteration: ids.runIteration,
+    project_id: ids.projectId,
+    dataset_id: ids.datasetId,
+    dataset_record_id: ids.datasetRecordId,
+    dataset_name: ids.datasetName,
+    experiment_name: ids.experimentName,
+  })
+
   return {
     span_id: row.spanId,
     trace_id: row.traceId,
@@ -43,13 +58,7 @@ function toSpan (row, metadata, ids, spanName, userTags) {
     duration: row.durationNs,
     status: row.isError ? 'error' : 'ok',
     meta,
-    tags: buildTags(userTags, {
-      experiment_id: ids.experimentId,
-      run_id: ids.runId,
-      run_iteration: ids.runIteration,
-      dataset_id: ids.datasetId,
-      dataset_record_id: ids.datasetRecordId,
-    }),
+    tags,
   }
 }
 
@@ -152,6 +161,8 @@ class Experiment {
     this.#evaluators = normalizeEvaluators(options.evaluators, 'row')
     this.#summaryEvaluators = normalizeEvaluators(options.summaryEvaluators, 'summary')
     this.#config = { ...options.config }
+    const filterTags = this.#dataset.filterTags?.() ?? []
+    if (filterTags.length > 0) this.#config.filtered_record_tags = filterTags
     this.#tags = { ...options.tags }
     this.#metadata = { ...options.metadata }
     this.#projectId = null
@@ -454,9 +465,11 @@ class Experiment {
             projectId,
             datasetId,
             datasetRecordId,
+            datasetName: this.#dataset.name(),
+            experimentName: this.#name,
             runId,
             runIteration,
-          }, this.#task.name || this.#name, this.#tags))
+          }, this.#task.name || this.#name, this.#tags, record.tags))
         }
       }
 
@@ -521,7 +534,7 @@ class Experiment {
       dataset_name: this.#dataset.name(),
       experiment_name: this.#name,
     }
-    const tags = mergeTags(this.#tags, autoTags)
+    const tags = mergeTags(this.#tags, { ...recordTagsToObject(record.tags), ...autoTags })
 
     const execute = () => this.#runWithRetries(
       () => this.#task(record.input, this.#config, record.metadata),

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

@@ -4,6 +4,7 @@ const log = require('../../log')
 const { ExperimentsClient } = require('./client')
 const { Dataset, DatasetRecord } = require('./dataset')
 const { Experiment, ExternalExperiment } = require('./experiment')
+const { validateTagsList } = require('./util')
 const NoopExperiments = require('./noop')
 
 // Poll `attempt` with exponential backoff until it returns true or the time
@@ -66,7 +67,9 @@ class Experiments {
           if (recordIds.has(record.id)) throw new Error(`Duplicate record id '${record.id}'`)
           recordIds.add(record.id)
         }
-        dataset.addRecord(new DatasetRecord(record.inputData, record.expectedOutput, record.metadata, record.id))
+        dataset.addRecord(
+          new DatasetRecord(record.inputData, record.expectedOutput, record.metadata, record.id, record.tags)
+        )
       }
     }
     return dataset
@@ -74,14 +77,16 @@ class Experiments {
 
   // Pull an existing dataset by name (with its records). Polls with exponential
   // backoff to absorb read-after-write lag; pass `expectedRecordCount` to also
-  // wait until that many records are readable.
+  // wait until that many records are readable. Pass `tags` to filter records by
+  // dataset record tags.
   async pullDataset (name, options = {}) {
-    const { expectedRecordCount, maxWaitMs = 30_000, version } = options
+    const { expectedRecordCount, maxWaitMs = 30_000, tags, version } = options
+    const filterTags = validateTagsList(tags)
     const projectId = await this.#client.ensureProjectId()
 
     let pulledDataset = null
     let records = []
-    let datasetVersion = version ?? null
+    const datasetVersion = version ?? null
     let latestVersion = null
     let lastError = ''
 
@@ -93,7 +98,6 @@ class Experiments {
             if (dataset.name() === name) {
               pulledDataset = dataset
               latestVersion = dataset.latestVersion()
-              datasetVersion = version ?? latestVersion
               break
             }
           }
@@ -107,6 +111,7 @@ class Experiments {
           // eslint-disable-next-line no-await-in-loop
           const page = await this.#client.listDatasetRecords(projectId, pulledDataset.id(), {
             cursor,
+            tags: filterTags,
             version: datasetVersion,
           })
           for (const record of page.records) recs.push(record)
@@ -139,6 +144,12 @@ class Experiments {
       )
     }
 
+    for (const record of records) {
+      if (record.id === null || record.id === undefined || record.id === '') {
+        throw new Error(`Failed to pull dataset '${name}': backend returned a record without an id`)
+      }
+    }
+
     return Dataset.fromExisting(
       this.#client,
       name,
@@ -147,7 +158,8 @@ class Experiments {
       projectId,
       records,
       datasetVersion,
-      latestVersion
+      latestVersion,
+      filterTags
     )
   }
 

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

@@ -11,20 +11,29 @@ class NoopDataset {
   #name
   #description
   #records
+  #filterTags
 
   constructor (name = '', options = {}) {
     this.#name = name
     this.#description = typeof options === 'string' ? options : (options.description ?? '')
+    this.#filterTags = typeof options === 'string' ? [] : [...(options.filterTags ?? [])]
     this.#records = (typeof options === 'string' ? [] : (options.records ?? [])).map(record => ({
       id: record.id ?? null,
       input: record.inputData,
       expectedOutput: record.expectedOutput ?? null,
       metadata: record.metadata ?? {},
+      tags: [...(record.tags ?? [])],
     }))
   }
 
-  addRecord (input, expectedOutput, metadata) {
-    this.#records.push({ id: null, input, expectedOutput: expectedOutput ?? null, metadata: metadata ?? {} })
+  addRecord (input, expectedOutput, metadata, tags) {
+    this.#records.push({
+      id: null,
+      input,
+      expectedOutput: expectedOutput ?? null,
+      metadata: metadata ?? {},
+      tags: [...(tags ?? [])],
+    })
     return this
   }
 
@@ -46,6 +55,49 @@ class NoopDataset {
     return Promise.resolve({ pushedCount: 0, totalCount: 0 })
   }
 
+  /**
+   * Add tags to a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {NoopDataset} This dataset for chaining.
+   */
+  addTags (index, tags) {
+    const record = this.#records[index]
+    if (!record) return this
+    const normalizedTags = tags ?? []
+    record.tags = [...new Set([...(record.tags ?? []), ...normalizedTags])].sort()
+    return this
+  }
+
+  /**
+   * Remove tags from a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {NoopDataset} This dataset for chaining.
+   */
+  removeTags (index, tags) {
+    const record = this.#records[index]
+    if (!record) return this
+    const normalizedTags = tags ?? []
+    const removed = new Set(normalizedTags)
+    record.tags = (record.tags ?? []).filter(tag => !removed.has(tag)).sort()
+    return this
+  }
+
+  /**
+   * Replace all tags on a dataset record.
+   * @param {number} index Dataset record index.
+   * @param {string[]} tags Tags in key:value format.
+   * @returns {NoopDataset} This dataset for chaining.
+   */
+  replaceTags (index, tags) {
+    const record = this.#records[index]
+    if (!record) return this
+    const normalizedTags = tags ?? []
+    record.tags = [...normalizedTags]
+    return this
+  }
+
   name () {
     return this.#name
   }
@@ -70,6 +122,14 @@ class NoopDataset {
     return null
   }
 
+  /**
+   * Return the tags used to filter this dataset.
+   * @returns {string[]} Dataset record filter tags.
+   */
+  filterTags () {
+    return [...this.#filterTags]
+  }
+
   records () {
     return [...this.#records]
   }
@@ -156,9 +216,9 @@ class NoopExperiments {
     return new NoopDataset(name, options)
   }
 
-  pullDataset (name) {
+  pullDataset (name, options = {}) {
     this.#warn()
-    return Promise.resolve(new NoopDataset(name))
+    return Promise.resolve(new NoopDataset(name, { filterTags: options.tags }))
   }
 
   experiment (options = {}) {