dd-trace 6.10.0 → 6.12.0

package/packages/dd-trace/src/exporters/common/buffering-exporter.js

@@ -24,23 +24,25 @@ class BufferingExporter {
     this._export(trace)
   }
 
-  _export (payload, writer = this._writer, timerKey = '_timer') {
+  _export (payload, writer = this._writer, timerKey = '_timer', deferImmediateFlush = false) {
     if (this._config.isCiVisibility) {
       incrementCountMetric(TELEMETRY_EVENTS_ENQUEUED_FOR_SERIALIZATION, {}, payload.length)
     }
-    writer.append(payload)
+    const appended = writer.append(payload)
 
     const { flushInterval } = this._config
 
-    if (flushInterval === 0) {
+    if (flushInterval === 0 && !deferImmediateFlush) {
       writer.flush()
-    } else if (this[timerKey] === undefined) {
+    } else if (flushInterval !== 0 && this[timerKey] === undefined) {
       this[timerKey] = setTimeout(() => {
         writer.flush()
         this[timerKey] = undefined
       }, flushInterval)
       this[timerKey].unref?.()
     }
+
+    return appended
   }
 
   getUncodedTraces () {

package/packages/dd-trace/src/exporters/common/request.js

@@ -93,7 +93,10 @@ function request (data, options, callback) {
 
   docker.inject(options.headers)
 
-  options.agent = isSecure ? httpsAgent : httpAgent
+  const connectionOptions = {
+    ...options,
+    agent: options.agent ?? (isSecure ? httpsAgent : httpAgent),
+  }
 
   /**
    * @param {import('node:http').IncomingMessage} res
@@ -218,7 +221,7 @@ function request (data, options, callback) {
         }
       }
 
-      const req = client.request(options, (res) => onResponse(res, complete, handleError))
+      const req = client.request(connectionOptions, (res) => onResponse(res, complete, handleError))
 
       req.once('close', finalize)
       req.once('timeout', finalize)

package/packages/dd-trace/src/exporters/common/writer.js

@@ -17,10 +17,10 @@ class Writer {
 
   #isFirstFlush = true
 
-  flush (done = () => {}) {
+  flush (done = () => {}, options) {
     const count = this._encoder.count()
 
-    if (!request.writable) {
+    if (!request.writable && options?.deadline === undefined) {
       this._encoder.reset()
       done()
     } else if (count > 0) {
@@ -44,23 +44,28 @@ class Writer {
         done()
         return
       }
-      this._sendPayload(payload, count, done)
+      if (options === undefined) {
+        this._sendPayload(payload, count, done)
+      } else {
+        this._sendPayload(payload, count, done, options)
+      }
     } else {
       done()
     }
   }
 
-  append (payload) {
-    if (!request.writable) {
+  append (payload, options) {
+    if (!request.writable && options?.deadline === undefined) {
       // eslint-disable-next-line eslint-rules/eslint-log-printf-style
       log.debug(() => `Maximum number of active requests reached. Payload discarded: ${safeJSONStringify(payload)}`)
-      return
+      return false
     }
 
     // eslint-disable-next-line eslint-rules/eslint-log-printf-style
     log.debug(() => `Encoding payload: ${safeJSONStringify(payload)}`)
 
     this._encode(payload)
+    return true
   }
 
   _encode (payload) {

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,68 @@ 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,
+    attrs.tags ?? []
+  )
+}
+
+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 +95,7 @@ class ExperimentsClient {
     this.#site = site
     this.#projectName = projectName
     this.#timeout = timeout
+    this.apiBase = `https://${apiHost(this.#site)}`
     this.#cachedProjectId = null
   }
 
@@ -61,7 +121,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 +152,110 @@ 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 datasetMutationResultFromResources(resources)
+  }
+
+  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))
+    if (Array.isArray(options.tags)) {
+      for (const tag of options.tags) query.append('filter[tags]', tag)
+    }
+    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 +264,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
   }
 }