dd-trace 6.6.0 → 6.8.0

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

@@ -3,35 +3,19 @@
 const id = require('../../id')
 
 const { API_BASE_PATH } = require('./client')
-const { Row, ExperimentResult } = require('./result')
-
-// Mirrors dd-trace-py's _generate_metric_from_evaluation: plain objects are
-// json, everything else falls through to the lowercased categorical fallback.
-function inferMetricType (value) {
-  if (typeof value === 'boolean') return 'boolean'
-  if (typeof value === 'number' && Number.isFinite(value)) return 'score'
-  if (value !== null && typeof value === 'object' && !Array.isArray(value)) return 'json'
-  return 'categorical'
-}
-
-function stringify (value) {
-  if (value === null || value === undefined) return ''
-  if (typeof value === 'string') return value.toLowerCase()
-  if (typeof value === 'object') return JSON.stringify(value).toLowerCase()
-  return String(value).toLowerCase()
-}
-
-// Build the tag list, letting auto tags win over user tags on key conflict.
-function buildTags (userTags, autoTags) {
-  const tags = new Map()
-  for (const [k, v] of Object.entries(userTags ?? {})) {
-    tags.set(k, `${k}:${v}`)
-  }
-  for (const [k, v] of Object.entries(autoTags)) {
-    tags.set(k, `${k}:${v}`)
-  }
-  return [...tags.values()]
-}
+const { Row, ExperimentResult, ExperimentRun } = require('./result')
+const {
+  buildExperimentTagObject,
+  buildSpanMetadata,
+  buildTags,
+  hasEntries,
+  inferMetricType,
+  normalizeEvaluators,
+  normalizeJsonMetricValue,
+  sleep,
+  stringify,
+  validateEvaluatorName,
+} = require('./util')
 
 // One span per experiment row (LLM Obs experiment span wire format).
 function toSpan (row, metadata, ids, spanName, userTags) {
@@ -40,7 +24,7 @@ function toSpan (row, metadata, ids, spanName, userTags) {
     output: row.output ?? null,
     expected_output: row.expectedOutput ?? null,
   }
-  if (metadata && Object.keys(metadata).length > 0) {
+  if (hasEntries(metadata)) {
     meta.metadata = metadata
   }
   if (row.isError) {
@@ -59,19 +43,26 @@ function toSpan (row, metadata, ids, spanName, userTags) {
     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,
     }),
   }
 }
 
-// One metric per evaluator per row.
-function toMetric (label, value, errorMessage, spanId, timestampMs, experimentId, userTags) {
+// One metric per evaluator per row or summary evaluator.
+function toMetric (
+  label, value, errorMessage, spanId, traceId, timestampMs, experimentId, userTags, source = 'custom'
+) {
   const metric = {
+    metric_source: source,
     label,
     span_id: spanId,
+    trace_id: traceId,
     timestamp_ms: timestampMs,
     tags: buildTags(userTags, { experiment_id: experimentId }),
+    experiment_id: experimentId,
   }
 
   if (errorMessage !== null) {
@@ -84,35 +75,50 @@ function toMetric (label, value, errorMessage, spanId, timestampMs, experimentId
   metric.metric_type = type
   if (type === 'boolean') metric.boolean_value = value
   else if (type === 'score') metric.score_value = value
-  else if (type === 'json') metric.json_value = value
+  else if (type === 'json') metric.json_value = normalizeJsonMetricValue(value)
   else metric.categorical_value = stringify(value)
   return metric
 }
 
+function createFallbackSpanContext (startNs) {
+  // Root-span id generation, same convention as opentracing/span.js: a single
+  // random 64-bit id for the span, reused as the trace id's low 64 bits with a
+  // start-time-derived high 64 bits (like the `_dd.p.tid` 128-bit trace id tag).
+  const spanIdentifier = id()
+  const traceIdHigh = Math.floor(startNs / 1e9).toString(16).padStart(8, '0').padEnd(16, '0')
+  const spanId = spanIdentifier.toString(16).padStart(16, '0')
+  const traceId = spanIdentifier.toTraceIdHex(traceIdHigh).padStart(32, '0')
+  return { spanId, traceId }
+}
+
 // Builder + run() orchestration: runs rows sequentially, emits one root span
-// per dataset row, and posts all spans + metrics in a single events call.
+// per dataset row, and posts spans + metrics to the experiments events API.
 class Experiment {
   #client
+  #llmobs
   #name
   #description
   #dataset
   #task
   #evaluators
+  #summaryEvaluators
   #config
   #tags
   #experimentId
 
-  constructor (client, options = {}) {
+  constructor (client, options = {}, llmobs) {
     if (!options.name) throw new Error('Experiment name is required')
     if (!options.dataset) throw new Error('Experiment dataset is required')
     if (typeof options.task !== 'function') throw new Error('Experiment task is required')
 
     this.#client = client
+    this.#llmobs = llmobs
     this.#name = options.name
     this.#description = options.description ?? ''
     this.#dataset = options.dataset
     this.#task = options.task
-    this.#evaluators = new Map(Object.entries(options.evaluators ?? {}))
+    this.#evaluators = normalizeEvaluators(options.evaluators, 'row')
+    this.#summaryEvaluators = normalizeEvaluators(options.summaryEvaluators, 'summary')
     this.#config = { ...options.config }
     this.#tags = { ...options.tags }
     this.#experimentId = null
@@ -131,7 +137,16 @@ class Experiment {
     return `${this.#client.appBase}/llm/experiments/${this.#experimentId}`
   }
 
-  async run () {
+  async run (options = {}) {
+    const {
+      maxRetries = 0,
+      retryDelay = (attempt) => 100 * (attempt + 1),
+      throwOnErrors = false,
+    } = options
+
+    if (maxRetries < 0) throw new Error('maxRetries must be >= 0')
+    if (typeof retryDelay !== 'function') throw new TypeError('retryDelay must be a function')
+
     const projectId = await this.#client.ensureProjectId()
 
     await this.#dataset.ensureCreatedAndPushed(projectId)
@@ -148,7 +163,11 @@ class Experiment {
       dataset_id: datasetId,
       description: this.#description,
       ensure_unique: true,
+      run_count: 1,
+      metadata: { tags: buildTags(this.#tags, {}) },
     }
+    const datasetVersion = this.#dataset.version()
+    if (datasetVersion !== null) attributes.dataset_version = datasetVersion
     if (Object.keys(this.#config).length > 0) attributes.config = this.#config
 
     let created
@@ -161,6 +180,8 @@ class Experiment {
     }
     this.#experimentId = created?.data?.id ?? null
     const experimentId = this.#experimentId
+    const runId = id().toString(16).padStart(16, '0')
+    const runIteration = 0
 
     try {
       const records = this.#dataset.records()
@@ -168,83 +189,87 @@ class Experiment {
       const rows = []
       const spans = []
       const metrics = []
+      const evaluatorResults = {}
+      const usesLLMObsTrace = Boolean(this.#llmobs?.enabled)
       let hasRowError = false
 
       for (let i = 0; i < records.length; i++) {
         const record = records[i]
         const datasetRecordId = i < recordIds.length ? recordIds[i] : ''
-        const startNs = Date.now() * 1e6
-        const startHr = process.hrtime.bigint()
-
-        // Root-span id generation, same convention as opentracing/span.js: a single
-        // random 64-bit id for the span, reused as the trace id's low 64 bits with a
-        // start-time-derived high 64 bits (like the `_dd.p.tid` 128-bit trace id tag).
-        const spanIdentifier = id()
-        const traceIdHigh = Math.floor(startNs / 1e9).toString(16).padStart(8, '0').padEnd(16, '0')
-        const spanId = spanIdentifier.toString(16).padStart(16, '0')
-        const traceId = spanIdentifier.toTraceIdHex(traceIdHigh).padStart(32, '0')
-
-        let output = null
-        let errorType = null
-        let errorMessage = null
-        try {
-          // Rows run one at a time by design.
-          // eslint-disable-next-line no-await-in-loop
-          output = await this.#task(record.input, this.#config)
-        } catch (err) {
-          errorType = err.name || 'Error'
-          errorMessage = err.message ?? String(err)
-        }
+        // Rows currently run sequentially by design; jobs/concurrency is a P1 follow-up.
+        // eslint-disable-next-line no-await-in-loop
+        const row = await this.#processRecord({
+          index: i,
+          record,
+          datasetRecordId,
+          projectId,
+          datasetId,
+          experimentId,
+          runId,
+          runIteration,
+          maxRetries,
+          retryDelay,
+          throwOnErrors,
+        })
 
-        const durationNs = Number(process.hrtime.bigint() - startHr)
-        const evaluations = {}
-        const evaluationErrors = {}
         const timestampMs = Date.now()
-
         for (const [label, evaluator] of this.#evaluators) {
-          if (errorType !== null) {
+          if (!evaluatorResults[label]) evaluatorResults[label] = []
+          if (row.isError) {
             const msg = 'task error; evaluation skipped'
-            evaluationErrors[label] = msg
-            metrics.push(toMetric(label, null, msg, spanId, timestampMs, experimentId, this.#tags))
+            row.evaluationErrors[label] = msg
+            evaluatorResults[label].push(null)
+            metrics.push(toMetric(label, null, msg, row.spanId, row.traceId, timestampMs, experimentId, this.#tags))
             continue
           }
           try {
             // eslint-disable-next-line no-await-in-loop
-            const value = await evaluator(record.input, output, record.expectedOutput)
-            evaluations[label] = value
-            metrics.push(toMetric(label, value, null, spanId, timestampMs, experimentId, this.#tags))
+            const value = await this.#runWithRetries(
+              () => evaluator(record.input, row.output, record.expectedOutput),
+              maxRetries,
+              retryDelay
+            )
+            row.evaluations[label] = value
+            evaluatorResults[label].push(value)
+            metrics.push(toMetric(label, value, null, row.spanId, row.traceId, timestampMs, experimentId, this.#tags))
           } catch (err) {
+            if (throwOnErrors) throw err
             const msg = err.message ?? String(err)
-            evaluationErrors[label] = msg
-            metrics.push(toMetric(label, null, msg, spanId, timestampMs, experimentId, this.#tags))
+            row.evaluationErrors[label] = msg
+            evaluatorResults[label].push(null)
+            metrics.push(toMetric(label, null, msg, row.spanId, row.traceId, timestampMs, experimentId, this.#tags))
           }
         }
 
-        const row = new Row({
-          index: i,
-          spanId,
-          traceId,
-          startNs,
-          durationNs,
-          input: record.input,
-          output,
-          expectedOutput: record.expectedOutput,
-          errorType,
-          errorMessage,
-          evaluations,
-          evaluationErrors,
-        })
         rows.push(row)
-        if (row.isError || Object.keys(evaluationErrors).length > 0) hasRowError = true
-        spans.push(toSpan(row, record.metadata, {
-          experimentId,
-          projectId,
-          datasetId,
-          datasetRecordId,
-        }, this.#name, this.#tags))
+        if (row.isError || hasEntries(row.evaluationErrors)) hasRowError = true
+        if (!usesLLMObsTrace) {
+          spans.push(toSpan(row, record.metadata, {
+            experimentId,
+            projectId,
+            datasetId,
+            datasetRecordId,
+            runId,
+            runIteration,
+          }, this.#task.name || this.#name, this.#tags))
+        }
+      }
+
+      const summaryEvaluations = await this.#runSummaryEvaluators(rows, records, evaluatorResults, {
+        maxRetries,
+        retryDelay,
+        throwOnErrors,
+        experimentId,
+        metrics,
+      })
+      if (hasEntries(summaryEvaluations)) {
+        for (const value of Object.values(summaryEvaluations)) {
+          if (value?.error) hasRowError = true
+        }
       }
 
       await this.#postEvents(experimentId, spans, metrics)
+      this.#llmobs?.flush?.()
       // A row error doesn't abort the run, but the experiment didn't succeed cleanly.
       await this.#updateStatus(
         experimentId,
@@ -252,16 +277,185 @@ class Experiment {
         hasRowError ? 'one or more rows failed' : null
       )
 
-      return new ExperimentResult(experimentId, rows, this.url())
+      const run = new ExperimentRun({ runId, runIteration, rows, summaryEvaluations })
+      return new ExperimentResult(experimentId, rows, this.url(), [run], summaryEvaluations)
     } catch (err) {
       await this.#updateStatus(experimentId, 'failed', err.message ?? String(err))
       throw err
     }
   }
 
+  async #processRecord ({
+    index,
+    record,
+    datasetRecordId,
+    projectId,
+    datasetId,
+    experimentId,
+    runId,
+    runIteration,
+    maxRetries,
+    retryDelay,
+    throwOnErrors,
+  }) {
+    const startNs = Date.now() * 1e6
+    const startHr = process.hrtime.bigint()
+    const fallbackContext = createFallbackSpanContext(startNs)
+    let spanContext = fallbackContext
+    let output = null
+    let errorType = null
+    let errorMessage = null
+
+    const autoTags = {
+      experiment_id: experimentId,
+      run_id: runId,
+      run_iteration: runIteration,
+      project_id: projectId,
+      dataset_id: datasetId,
+      dataset_record_id: datasetRecordId,
+      dataset_name: this.#dataset.name(),
+      experiment_name: this.#name,
+    }
+    const tags = buildExperimentTagObject(this.#tags, autoTags)
+
+    const execute = () => this.#runWithRetries(
+      () => this.#task(record.input, this.#config, record.metadata),
+      maxRetries,
+      retryDelay
+    )
+
+    if (this.#llmobs?.enabled) {
+      try {
+        await this.#llmobs.trace({
+          kind: 'experiment',
+          name: this.#task.name || this.#name,
+        }, async (span) => {
+          spanContext = this.#llmobs.exportSpan(span)
+          try {
+            output = await execute()
+          } catch (err) {
+            this.#llmobs.annotate(span, {
+              inputData: record.input,
+              outputData: output,
+              metadata: buildSpanMetadata(record.metadata, this.#config),
+              tags,
+            })
+            throw err
+          }
+          this.#llmobs.annotate(span, {
+            inputData: record.input,
+            outputData: output,
+            metadata: buildSpanMetadata(record.metadata, this.#config),
+            tags,
+          })
+        })
+      } catch (err) {
+        if (throwOnErrors) throw err
+        errorType = err.name || 'Error'
+        errorMessage = err.message ?? String(err)
+      }
+    } else {
+      try {
+        output = await execute()
+      } catch (err) {
+        if (throwOnErrors) throw err
+        errorType = err.name || 'Error'
+        errorMessage = err.message ?? String(err)
+      }
+    }
+
+    const durationNs = Number(process.hrtime.bigint() - startHr)
+    return new Row({
+      index,
+      spanId: spanContext?.spanId ?? fallbackContext.spanId,
+      traceId: spanContext?.traceId ?? fallbackContext.traceId,
+      startNs,
+      durationNs,
+      input: record.input,
+      output,
+      expectedOutput: record.expectedOutput,
+      errorType,
+      errorMessage,
+      evaluations: {},
+      evaluationErrors: {},
+    })
+  }
+
+  async #runWithRetries (fn, maxRetries, retryDelay) {
+    let lastError
+    for (let attempt = 0; attempt <= maxRetries; attempt++) {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        return await fn()
+      } catch (err) {
+        lastError = err
+        if (attempt >= maxRetries) break
+        const delayMs = retryDelay(attempt)
+        if (typeof delayMs !== 'number' || delayMs < 0 || !Number.isFinite(delayMs)) {
+          throw new TypeError('retryDelay must return a non-negative finite number of milliseconds')
+        }
+        // eslint-disable-next-line no-await-in-loop
+        await sleep(delayMs)
+      }
+    }
+    throw lastError
+  }
+
+  async #runSummaryEvaluators (rows, records, evaluatorResults, options) {
+    if (this.#summaryEvaluators.length === 0) return {}
+
+    const inputs = rows.map(row => row.input)
+    const outputs = rows.map(row => row.output)
+    const expectedOutputs = rows.map(row => row.expectedOutput)
+    const metadata = records.map(record => buildSpanMetadata(record.metadata, this.#config))
+    const summaryEvaluations = {}
+    const timestampMs = Date.now()
+
+    for (const [label, evaluator] of this.#summaryEvaluators) {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        const value = await this.#runWithRetries(
+          () => evaluator(inputs, outputs, expectedOutputs, evaluatorResults, metadata),
+          options.maxRetries,
+          options.retryDelay
+        )
+        summaryEvaluations[label] = { value, error: null }
+        options.metrics.push(toMetric(
+          label,
+          value,
+          null,
+          '',
+          '',
+          timestampMs,
+          options.experimentId,
+          this.#tags,
+          'summary'
+        ))
+      } catch (err) {
+        if (options.throwOnErrors) throw err
+        const msg = err.message ?? String(err)
+        summaryEvaluations[label] = { value: null, error: msg }
+        options.metrics.push(toMetric(
+          label,
+          null,
+          msg,
+          '',
+          '',
+          timestampMs,
+          options.experimentId,
+          this.#tags,
+          'summary'
+        ))
+      }
+    }
+    return summaryEvaluations
+  }
+
   async #postEvents (experimentId, spans, metrics) {
+    const attributes = { metrics }
+    if (spans.length > 0) attributes.spans = spans
     await this.#client.request('POST', `${API_BASE_PATH}/experiments/${experimentId}/events`, {
-      data: { type: 'experiments', attributes: { spans, metrics } },
+      data: { type: 'experiments', attributes },
     })
   }
 
@@ -280,4 +474,4 @@ class Experiment {
   }
 }
 
-module.exports = { Experiment }
+module.exports = { Experiment, normalizeEvaluators, validateEvaluatorName }

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

@@ -26,9 +26,11 @@ async function retryWithBackoff (attempt, { maxTotalMs = 30_000, baseDelayMs = 2
 // experiments against the LLM Obs backend using the tracer's own config.
 class Experiments {
   #client
+  #llmobs
   #projectName
 
-  constructor (config) {
+  constructor (config, llmobs) {
+    this.#llmobs = llmobs
     this.#projectName = config.llmobs?.mlApp || config.service
     this.#client = new ExperimentsClient({
       apiKey: config.DD_API_KEY,
@@ -39,21 +41,38 @@ class Experiments {
   }
 
   // Create a local dataset buffer. Pushed remotely on first experiment run.
-  createDataset (name, description = '') {
-    return new Dataset(this.#client, name, description)
+  createDataset (name, descriptionOrOptions = '') {
+    const options = typeof descriptionOrOptions === 'string'
+      ? { description: descriptionOrOptions }
+      : (descriptionOrOptions ?? {})
+    const dataset = new Dataset(this.#client, name, options.description ?? '')
+    const recordIds = new Set()
+    for (const record of options.records ?? []) {
+      if (record.id !== undefined && (typeof record.id !== 'string' || record.id.length === 0)) {
+        throw new Error('record id must be a non-empty string')
+      }
+      if (record.id !== undefined) {
+        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))
+    }
+    return dataset
   }
 
   // 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.
   async pullDataset (name, options = {}) {
-    const { expectedRecordCount, maxWaitMs = 30_000 } = options
+    const { expectedRecordCount, maxWaitMs = 30_000, version } = options
     const projectId = await this.#client.ensureProjectId()
 
     let datasetId = null
     let description = ''
     let records = []
     let recordIds = []
+    let datasetVersion = version ?? null
+    let latestVersion = null
     let lastError = ''
 
     const succeeded = await retryWithBackoff(async () => {
@@ -67,6 +86,8 @@ class Experiments {
             if (item?.attributes?.name === name) {
               datasetId = String(item?.id ?? '')
               description = String(item?.attributes?.description ?? '')
+              latestVersion = item?.attributes?.current_version ?? null
+              datasetVersion = version ?? latestVersion
               break
             }
           }
@@ -78,16 +99,24 @@ class Experiments {
         let cursor = ''
         // Follow the meta.after / page[cursor] pagination until the last page.
         for (;;) {
-          const query = cursor ? `?page[cursor]=${encodeURIComponent(cursor)}` : ''
+          const query = new URLSearchParams()
+          if (cursor) query.set('page[cursor]', cursor)
+          if (datasetVersion !== null) query.set('filter[version]', String(datasetVersion))
           // eslint-disable-next-line no-await-in-loop
           const resp = await this.#client.request(
             'GET',
-            `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records${query}`
+            `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records?${query.toString()}`
           )
           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
@@ -119,18 +148,28 @@ class Experiments {
       )
     }
 
-    return Dataset.fromExisting(this.#client, name, description, datasetId, projectId, records, recordIds)
+    return Dataset.fromExisting(
+      this.#client,
+      name,
+      description,
+      datasetId,
+      projectId,
+      records,
+      recordIds,
+      datasetVersion,
+      latestVersion
+    )
   }
 
   // Build an experiment: { name, dataset, task, evaluators, description?, config?, tags? }.
   experiment (options) {
-    return new Experiment(this.#client, options)
+    return new Experiment(this.#client, options, this.#llmobs)
   }
 }
 
 // Factory used by the LLMObs SDK: returns a real Experiments instance when
 // enabled and credentialed, otherwise a no-op that explains what's missing.
-function createExperiments (config) {
+function createExperiments (config, llmobs) {
   if (!config.llmobs?.DD_LLMOBS_ENABLED) {
     return new NoopExperiments('LLM Observability is not enabled')
   }
@@ -146,7 +185,7 @@ function createExperiments (config) {
       'then retry'
     )
   }
-  return new Experiments(config)
+  return new Experiments(config, llmobs)
 }
 
 module.exports = { Experiments, createExperiments }