feat(llmobs): add experiments dataset operations (#9458)

* feat(llmobs): add experiments dataset operations

* feat(llmobs): add CSV dataset creation

* chore(llmobs): use csv parser for dataset import

* chore(llmobs): add csv parser license

* feat(llmobs): include dataset version in experiments

* fix(llmobs): address dataset review feedback

* create jsontype to replace any

* remove dataset operations from llmobs

* remove csv dataset func

* noops produce warn logs instead of error throwing

* use vcr testing instead

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

Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com>

* fix(llmobs): address experiments ci failures

* test(llmobs): cover experiments noop accessors

* Apply suggestion from @sabrenner

Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com>

* Apply suggestion from @sabrenner

Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com>

* fix(llmobs): simplify dataset records query route

* test(llmobs): remove handmade experiment cassettes

* test(llmobs): isolate experiment fetch mocks

---------

Co-authored-by: Sam Brenner <106700075+sabrenner@users.noreply.github.com>

Author: mehulsonowal

index.d.ts

@@ -3734,16 +3734,36 @@ declare namespace tracer {
       flush (): void
     }
 
+    /** JSON-serializable value accepted by LLMObs Experiments. */
+    type JSONType = string | number | boolean | null | JSONType[] | { [key: string]: JSONType }
+
     /**
      * A task run over each dataset record during an experiment.
      */
-    type ExperimentTask = (input: any, config: Record<string, any>) => any | Promise<any>
+    type ExperimentTask = (
+      input: JSONType,
+      config: Record<string, JSONType>
+    ) => JSONType | Promise<JSONType>
 
     /**
      * Scores a single task output. The return type selects the metric:
      * `boolean` -> boolean, `number` -> score, anything else -> categorical.
      */
-    type ExperimentEvaluator = (input: any, output: any, expectedOutput: any) => any | Promise<any>
+    type ExperimentEvaluator = (
+      input: JSONType,
+      output: JSONType,
+      expectedOutput: JSONType
+    ) => JSONType | Promise<JSONType>
+
+    interface CreateDatasetOptions {
+      description?: string
+      records?: Array<{
+        id?: string,
+        inputData: JSONType,
+        expectedOutput?: JSONType,
+        metadata?: Record<string, JSONType>
+      }>
+    }
 
     interface ExperimentOptions {
       name: string
@@ -3752,11 +3772,13 @@ declare namespace tracer {
       /** Evaluators keyed by metric label. */
       evaluators?: Record<string, ExperimentEvaluator>
       description?: string
-      config?: Record<string, any>
+      config?: Record<string, JSONType>
       tags?: Record<string, string>
     }
 
     interface PullDatasetOptions {
+      /** Dataset version to pull. Defaults to latest. */
+      version?: number
       /** Wait until at least this many records are readable (absorbs write lag). */
       expectedRecordCount?: number
       /** Maximum total time to wait, in ms. Default 30000. */
@@ -3769,13 +3791,13 @@ declare namespace tracer {
       traceId: string
       startNs: number
       durationNs: number
-      input: any
-      output: any
-      expectedOutput: any
+      input: JSONType
+      output: JSONType
+      expectedOutput: JSONType
       readonly isError: boolean
       errorType: string | null
       errorMessage: string | null
-      evaluations: Record<string, any>
+      evaluations: Record<string, JSONType>
       evaluationErrors: Record<string, string>
     }
 
@@ -3794,13 +3816,20 @@ declare namespace tracer {
     }
 
     interface Dataset {
-      addRecord (input: any, expectedOutput?: any, metadata?: Record<string, any>): Dataset
+      addRecord (input: JSONType, expectedOutput?: JSONType, metadata?: Record<string, JSONType>): Dataset
       /** Creates the dataset remotely if needed and pushes any unpushed records. */
       push (): Promise<DatasetPushResult>
       name (): string
       id (): string | null
       projectId (): string | null
-      records (): Array<{ input: any, expectedOutput: any, metadata: Record<string, any> }>
+      version (): number | null
+      latestVersion (): number | null
+      records (): Array<{
+        id: string | null,
+        input: JSONType,
+        expectedOutput: JSONType,
+        metadata: Record<string, JSONType>
+      }>
       /** Dashboard URL for the dataset, or null until pushed. */
       url (): string | null
     }
@@ -3815,6 +3844,7 @@ declare namespace tracer {
     interface Experiments {
       /** Create a local dataset buffer; pushed on the first experiment run. */
       createDataset (name: string, description?: string): Dataset
+      createDataset (name: string, options?: CreateDatasetOptions): Dataset
       /** Pull an existing dataset (with records) by name. */
       pullDataset (name: string, options?: PullDatasetOptions): Promise<Dataset>
       /** Build an experiment to run over a dataset. */

index.d.v5.ts

@@ -3931,16 +3931,36 @@ declare namespace tracer {
       flush (): void
     }
 
+    /** JSON-serializable value accepted by LLMObs Experiments. */
+    type JSONType = string | number | boolean | null | JSONType[] | { [key: string]: JSONType }
+
     /**
      * A task run over each dataset record during an experiment.
      */
-    type ExperimentTask = (input: any, config: Record<string, any>) => any | Promise<any>
+    type ExperimentTask = (
+      input: JSONType,
+      config: Record<string, JSONType>
+    ) => JSONType | Promise<JSONType>
 
     /**
      * Scores a single task output. The return type selects the metric:
      * `boolean` -> boolean, `number` -> score, anything else -> categorical.
      */
-    type ExperimentEvaluator = (input: any, output: any, expectedOutput: any) => any | Promise<any>
+    type ExperimentEvaluator = (
+      input: JSONType,
+      output: JSONType,
+      expectedOutput: JSONType
+    ) => JSONType | Promise<JSONType>
+
+    interface CreateDatasetOptions {
+      description?: string
+      records?: Array<{
+        id?: string,
+        inputData: JSONType,
+        expectedOutput?: JSONType,
+        metadata?: Record<string, JSONType>
+      }>
+    }
 
     interface ExperimentOptions {
       name: string
@@ -3949,11 +3969,13 @@ declare namespace tracer {
       /** Evaluators keyed by metric label. */
       evaluators?: Record<string, ExperimentEvaluator>
       description?: string
-      config?: Record<string, any>
+      config?: Record<string, JSONType>
       tags?: Record<string, string>
     }
 
     interface PullDatasetOptions {
+      /** Dataset version to pull. Defaults to latest. */
+      version?: number
       /** Wait until at least this many records are readable (absorbs write lag). */
       expectedRecordCount?: number
       /** Maximum total time to wait, in ms. Default 30000. */
@@ -3966,13 +3988,13 @@ declare namespace tracer {
       traceId: string
       startNs: number
       durationNs: number
-      input: any
-      output: any
-      expectedOutput: any
+      input: JSONType
+      output: JSONType
+      expectedOutput: JSONType
       readonly isError: boolean
       errorType: string | null
       errorMessage: string | null
-      evaluations: Record<string, any>
+      evaluations: Record<string, JSONType>
       evaluationErrors: Record<string, string>
     }
 
@@ -3991,13 +4013,20 @@ declare namespace tracer {
     }
 
     interface Dataset {
-      addRecord (input: any, expectedOutput?: any, metadata?: Record<string, any>): Dataset
+      addRecord (input: JSONType, expectedOutput?: JSONType, metadata?: Record<string, JSONType>): Dataset
       /** Creates the dataset remotely if needed and pushes any unpushed records. */
       push (): Promise<DatasetPushResult>
       name (): string
       id (): string | null
       projectId (): string | null
-      records (): Array<{ input: any, expectedOutput: any, metadata: Record<string, any> }>
+      version (): number | null
+      latestVersion (): number | null
+      records (): Array<{
+        id: string | null,
+        input: JSONType,
+        expectedOutput: JSONType,
+        metadata: Record<string, JSONType>
+      }>
       /** Dashboard URL for the dataset, or null until pushed. */
       url (): string | null
     }
@@ -4012,6 +4041,7 @@ declare namespace tracer {
     interface Experiments {
       /** Create a local dataset buffer; pushed on the first experiment run. */
       createDataset (name: string, description?: string): Dataset
+      createDataset (name: string, options?: CreateDatasetOptions): Dataset
       /** Pull an existing dataset (with records) by name. */
       pullDataset (name: string, options?: PullDatasetOptions): Promise<Dataset>
       /** Build an experiment to run over a dataset. */

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

@@ -14,9 +14,11 @@ function apiHost (site) {
 }
 
 // Web-app host for dashboard URLs. Single-level sites (datadoghq.com,
-// ddog-gov.com) are served from the `app.` subdomain; regional sites
-// (us3.datadoghq.com, ap1.datadoghq.com) are used as-is.
+// ddog-gov.com) are served from the `app.` subdomain; staging uses
+// dd.datad0g.com; regional sites (us3.datadoghq.com, ap1.datadoghq.com)
+// are used as-is.
 function appHost (site) {
+  if (site === 'datad0g.com') return 'dd.datad0g.com'
   return site.split('.').length === 2 ? `app.${site}` : site
 }
 

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

@@ -2,15 +2,35 @@
 
 const { API_BASE_PATH } = require('./client')
 
-// Immutable dataset record: { input, expectedOutput?, metadata? }.
+// 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 = {}) {
+  constructor (input, expectedOutput = null, metadata = {}, id = null) {
     this.input = input
     this.expectedOutput = expectedOutput ?? null
     this.metadata = metadata ?? {}
+    this.id = id ?? null
   }
 }
 
+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 {
@@ -22,6 +42,8 @@ class Dataset {
   #id
   #projectId
   #pushedCount
+  #version
+  #latestVersion
 
   constructor (client, name, description = '') {
     this.#client = client
@@ -32,16 +54,20 @@ class Dataset {
     this.#id = null
     this.#projectId = null
     this.#pushedCount = 0
+    this.#version = null
+    this.#latestVersion = null
   }
 
   // Build a Dataset that already exists remotely (used by pullDataset).
-  static fromExisting (client, name, description, id, projectId, records, recordIds) {
+  static fromExisting (client, name, description, id, projectId, records, recordIds, 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
     return dataset
   }
 
@@ -74,6 +100,14 @@ class Dataset {
     return this.#projectId
   }
 
+  version () {
+    return this.#version
+  }
+
+  latestVersion () {
+    return this.#latestVersion
+  }
+
   // Dashboard URL for this dataset, or null until pushed/pulled.
   url () {
     if (this.#id === null) return null
@@ -100,14 +134,22 @@ 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
     }
 
     if (this.#pushedCount >= this.#records.length) return { pushedCount: 0, totalCount: 0 }
 
     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
+      }
       if (rec.expectedOutput !== null && rec.expectedOutput !== undefined) {
         out.expected_output = rec.expectedOutput
       }
@@ -128,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/experiment.js

@@ -149,6 +149,8 @@ class Experiment {
       description: this.#description,
       ensure_unique: true,
     }
+    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