newrelic/lib/llm-events/chat-completion-message.js

New Relic agent

/*
 * Copyright 2026 New Relic Corporation. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

'use strict'

const LlmEvent = require('./event-base')

/**
 * An event that corresponds to each message (sent and received)
 * from a chat completion call including those created by the user,
 * assistant, and the system.
 *
 * @augments LlmEvent
 * @property {string} completion_id ID of the `LlmChatCompletionSummary` event
 *  that this message event is connected to
 * @property {string} content Content of the message
 * @property {string} id ID in the format `response_id`-`sequence`,
 *  or a UUID generated by the agent if no response ID is returned by the LLM
 * @property {boolean|undefined} is_response `true` if a message is the result
 *  of a chat completion and not an input message, `undefined` in `false` cases
 * @property {string} role Role of the message creator (e.g. `user`, `assistant`, `tool`)
 * @property {number} sequence Index (beginning at 0) associated with
 *  each message including the prompt and responses
 * @property {number|undefined} timestamp Timestamp captured at the time of the LLM
 *  request with millisecond precision, `undefined` if not a request
 */
module.exports = class LlmChatCompletionMessage extends LlmEvent {
  /**
   * @param {object} params Constructor parameters
   * @param {Agent} params.agent New Relic agent instance
   * @param {TraceSegment} params.segment Current segment
   * @param {Transaction} params.transaction Current and active transaction
   * @param {string} params.vendor Lowercase name of vendor (e.g. 'openai')
   * @param {string} params.requestId ID associated with the request -
   *    typically available in response headers
   * @param {string} params.responseId ID associated with the response, used to create `this.id`
   * @param {string} params.responseModel Model name returned in the response
   * @param {number} params.sequence Index (beginning at 0) associated with
   *    each message including the prompt and responses
   * @param {string} params.content Content of the message
   * @param {string} [params.role] Role of the message creator (e.g. `user`, `assistant`, `tool`)
   * @param {string} params.completionId ID of the `LlmChatCompletionSummary` event that
   *    this message event is connected to
   * @param {boolean} params.isResponse `true` if a message is the result of a chat
   *    completion and not an input message - omitted in `false` cases
   */
  constructor({ agent, segment, transaction, vendor, requestId, responseId, responseModel, sequence, content, role, completionId, isResponse }) {
    super({ agent, segment, transaction, vendor, responseModel, requestId })

    this.completion_id = completionId
    this.sequence = sequence
    if (isResponse) this.is_response = isResponse

    if (role) {
      this.role = role
    } else {
      // If the role attribute is not available, a value of user MUST be sent for
      // requests and a value of assistant MUST be sent for responses.
      if (isResponse) {
        this.role = 'assistant'
      } else if (sequence === 0) {
        // We can assume the first message in the sequence is the request message.
        this.role = 'user'
      }
    }

    if (isResponse !== true) {
      // Only include for input/request messages
      this.timestamp = segment.timer.start
    }

    if (responseId) {
      // A UUID is generated for `id` in super constructor,
      // but use this id format if responseId exists
      this.id = `${responseId}-${sequence}`
    }

    if (agent.config.ai_monitoring.record_content.enabled === true) {
      this.content = content
    }
  }

  /**
   * Sets `token_count` to 0 on the LlmChatCompletionMessage if both prompt and completion tokens are greater than zero.
   * This is because the spec states that if token counts are set, then we should set token_count to 0 to indicate
   * that the token calculation does not have to occur in the ingest pipeline.
   * @param {object} params to the function
   * @param {object} params.promptTokens value of prompt token count
   * @param {object} params.completionTokens value of completion(s) token count
   */
  setTokenInCompletionMessage({ promptTokens, completionTokens }) {
    if (this.isValidTokenCount(promptTokens) && this.isValidTokenCount(completionTokens)) {
      this.token_count = 0
    }
  }

  /**
   * Calculates prompt and completion token counts using the provided callback and models.
   * If both counts are valid, sets this.token_count to 0.
   *
   * @param {object} options The params object.
   * @param {Function} options.tokenCB The token counting callback function.
   * @param {string} options.reqModel The model used for the prompt.
   * @param {string} options.resModel The model used for the completion.
   * @param {string} options.promptContent The prompt content to count tokens for.
   * @param {string} options.completionContent The completion content to count tokens for.
   * @returns {void}
   */
  setTokenFromCallback({ tokenCB, reqModel, resModel, promptContent, completionContent }) {
    const promptToken = this.calculateCallbackTokens(tokenCB, reqModel, promptContent)
    const completionToken = this.calculateCallbackTokens(tokenCB, resModel, completionContent)

    this.setTokenInCompletionMessage({ promptTokens: promptToken, completionTokens: completionToken })
  }
}