chore: Removed unused attributes on LlmChatCompletionSummary and LlmE…

…mbedding (newrelic#278)

lib/llm/bedrock-response.js

@@ -122,25 +122,6 @@ class BedrockResponse {
  return this.#id
  }
 
- /**
- * The number of tokens present in the prompt as determined by the remote
- * API.
- *
- * @returns {number|undefined}
- */
- get inputTokenCount() {
- return this.#tokenCount('x-amzn-bedrock-input-token-count')
- }
-
- /**
- * The number of tokens in the LLM response as determined by the remote API.
- *
- * @returns {number|undefined}
- */
- get outputTokenCount() {
- return this.#tokenCount('x-amzn-bedrock-output-token-count')
- }
-
  /**
  * UUID assigned to the initial request as returned by the API.
  *

lib/llm/chat-completion-summary.js

@@ -28,17 +28,11 @@ class LlmChatCompletionSummary extends LlmEvent {
  this.duration = segment.getDurationInMillis()
  this['request.max_tokens'] = this.bedrockCommand.maxTokens
 
- const utt = 'response.usage.total_tokens'
  const nm = 'response.number_of_messages'
- const upt = 'response.usage.prompt_tokens'
- const uct = 'response.usage.completion_tokens'
  const cfr = 'response.choices.finish_reason'
  const rt = 'request.temperature'
 
  const cmd = this.bedrockCommand
- this[uct] = this.bedrockResponse.outputTokenCount
- this[upt] = this.bedrockResponse.inputTokenCount
- this[utt] = this[upt] + this[uct]
  this[cfr] = this.bedrockResponse.finishReason
  this[rt] = cmd.temperature
  this[nm] = 1 + this.bedrockResponse.completions.length

lib/llm/embedding.js

@@ -20,15 +20,16 @@ class LlmEmbedding extends LlmEvent {
  constructor(params = defaultParams) {
  super(params)
  const { agent } = params
+ const tokenCb = agent?.llm?.tokenCountCallback
 
  this.input = agent.config?.ai_monitoring?.record_content?.enabled
  ? this.bedrockCommand.prompt
  : undefined
  this.error = params.isError
  this.duration = params.segment.getDurationInMillis()
- this['response.usage.total_tokens'] = this.bedrockResponse.inputTokenCount
- this['response.usage.prompt_tokens'] = this.bedrockResponse.inputTokenCount
- this.token_count = this.bedrockResponse.inputTokenCount
+ if (typeof tokenCb === 'function') {
+  this.token_count = tokenCb(this.bedrockCommand.modelId, this.bedrockCommand.prompt)
+ }
  }
 }
 

lib/llm/event.js

@@ -56,7 +56,6 @@ class LlmEvent {
  this.ingest_source = 'Node'
  this.appName = agent.config.applications()[0]
  this.span_id = segment.id
- this.transaction_id = segment.transaction.id
  this.trace_id = segment.transaction.traceId
  this.request_id = this.bedrockResponse.requestId
  this.metadata = agent

lib/llm/stream-handler.js

@@ -181,10 +181,7 @@ class StreamHandler {
  return
  }
 
- const metrics = parsedEvent['amazon-bedrock-invocationMetrics']
  this.response.response.headers = {
- 'x-amzn-bedrock-input-token-count': metrics.inputTokenCount,
- 'x-amzn-bedrock-output-token-count': metrics.outputTokenCount,
  'x-amzn-requestid': this.passThroughParams.response.response.headers['x-amzn-requestid']
  }
  delete parsedEvent['amazon-bedrock-invocationMetrics']

tests/unit/llm/bedrock-response.tap.js

@@ -58,9 +58,7 @@ tap.beforeEach((t) => {
  statusCode: 200,
  headers: {
  'x-amzn-requestid': 'aws-request-1',
- 'x-foo': 'foo',
- ['x-amzn-bedrock-input-token-count']: 25,
- ['x-amzn-bedrock-output-token-count']: 25
+ 'x-foo': 'foo'
  }
  },
  output: {
@@ -98,8 +96,6 @@ tap.test('non-conforming response is handled gracefully', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, undefined)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, undefined)
- t.equal(res.outputTokenCount, undefined)
  t.equal(res.requestId, undefined)
  t.equal(res.statusCode, 200)
 })
@@ -111,8 +107,6 @@ tap.test('ai21 malformed responses work', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -125,8 +119,6 @@ tap.test('ai21 complete responses work', async (t) => {
  t.equal(res.finishReason, 'done')
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, 'ai21-response-1')
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -138,8 +130,6 @@ tap.test('claude malformed responses work', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -152,8 +142,6 @@ tap.test('claude complete responses work', async (t) => {
  t.equal(res.finishReason, 'done')
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -165,8 +153,6 @@ tap.test('cohere malformed responses work', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -179,8 +165,6 @@ tap.test('cohere complete responses work', async (t) => {
  t.equal(res.finishReason, 'done')
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, 'cohere-response-1')
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -192,8 +176,6 @@ tap.test('llama2 malformed responses work', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -206,8 +188,6 @@ tap.test('llama2 complete responses work', async (t) => {
  t.equal(res.finishReason, 'done')
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -219,8 +199,6 @@ tap.test('titan malformed responses work', async (t) => {
  t.equal(res.finishReason, undefined)
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })
@@ -233,8 +211,6 @@ tap.test('titan complete responses work', async (t) => {
  t.equal(res.finishReason, 'done')
  t.same(res.headers, t.context.response.response.headers)
  t.equal(res.id, undefined)
- t.equal(res.inputTokenCount, 25)
- t.equal(res.outputTokenCount, 25)
  t.equal(res.requestId, 'aws-request-1')
  t.equal(res.statusCode, 200)
 })

tests/unit/llm/chat-completion-message.tap.js

@@ -13,6 +13,7 @@ const LlmChatCompletionMessage = require('../../../lib/llm/chat-completion-messa
 
 tap.beforeEach((t) => {
  t.context.agent = {
+ llm: {},
  config: {
  applications() {
  return ['test-app']
@@ -85,6 +86,7 @@ tap.beforeEach((t) => {
 })
 
 tap.test('create creates a non-response instance', async (t) => {
+ t.context.agent.llm.tokenCountCallback = () => 3
  const event = new LlmChatCompletionMessage(t.context)
  t.equal(event.is_response, false)
  t.equal(event['llm.conversation_id'], 'conversation-1')
@@ -93,6 +95,7 @@ tap.test('create creates a non-response instance', async (t) => {
  t.equal(event.content, 'who are you')
  t.equal(event.role, 'user')
  t.match(event.id, /[\w-]{36}/)
+ t.equal(event.token_count, 3)
 })
 
 tap.test('create creates a titan response instance', async (t) => {

tests/unit/llm/chat-completion-summary.tap.js

@@ -70,9 +70,7 @@ tap.beforeEach((t) => {
  'x-amzn-request-id': 'aws-request-1'
  },
  finishReason: 'done',
- completions: ['completion-1'],
- inputTokenCount: 25,
- outputTokenCount: 25
+ completions: ['completion-1']
  }
 })
 
@@ -85,9 +83,6 @@ tap.test('creates a basic summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 0)
- t.equal(event['response.usage.prompt_tokens'], 0)
- t.equal(event['response.usage.completion_tokens'], 0)
  t.equal(event['response.number_of_messages'], 2)
 })
 
@@ -99,9 +94,6 @@ tap.test('creates an ai21 summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 50)
- t.equal(event['response.usage.prompt_tokens'], 25)
- t.equal(event['response.usage.completion_tokens'], 25)
  t.equal(event['response.number_of_messages'], 2)
 })
 
@@ -113,9 +105,6 @@ tap.test('creates an claude summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 50)
- t.equal(event['response.usage.prompt_tokens'], 25)
- t.equal(event['response.usage.completion_tokens'], 25)
  t.equal(event['response.number_of_messages'], 2)
 })
 
@@ -127,9 +116,6 @@ tap.test('creates a cohere summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 50)
- t.equal(event['response.usage.prompt_tokens'], 25)
- t.equal(event['response.usage.completion_tokens'], 25)
  t.equal(event['response.number_of_messages'], 2)
 })
 
@@ -141,9 +127,6 @@ tap.test('creates a llama2 summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 50)
- t.equal(event['response.usage.prompt_tokens'], 25)
- t.equal(event['response.usage.completion_tokens'], 25)
  t.equal(event['response.number_of_messages'], 2)
 })
 
@@ -155,8 +138,5 @@ tap.test('creates a titan summary', async (t) => {
  t.equal(event['request.max_tokens'], 25)
  t.equal(event['request.temperature'], 0.5)
  t.equal(event['response.choices.finish_reason'], 'done')
- t.equal(event['response.usage.total_tokens'], 50)
- t.equal(event['response.usage.prompt_tokens'], 25)
- t.equal(event['response.usage.completion_tokens'], 25)
  t.equal(event['response.number_of_messages'], 2)
 })

tests/unit/llm/embedding.tap.js

@@ -13,6 +13,7 @@ const LlmEmbedding = require('../../../lib/llm/embedding')
 
 tap.beforeEach((t) => {
  t.context.agent = {
+ llm: {},
  config: {
  applications() {
  return ['test-app']
@@ -49,13 +50,10 @@ tap.beforeEach((t) => {
  t.context.bedrockResponse = {
  headers: {
  'x-amzn-requestid': 'request-1'
- },
- get inputTokenCount() {
- return 8
  }
  }
  t.context.segment = {
- transaction: { id: '1', traceId: 'id' },
+ transaction: { traceId: 'id' },
  getDurationInMillis() {
  return 1.008
  }
@@ -66,9 +64,7 @@ tap.test('creates a basic embedding', async (t) => {
  const event = new LlmEmbedding(t.context)
  t.equal(event.input, 'who are you')
  t.equal(event.duration, 1.008)
- t.equal(event['response.usage.total_tokens'], 8)
- t.equal(event['response.usage.prompt_tokens'], 8)
- t.equal(event.token_count, 8)
+ t.equal(event.token_count, undefined)
 })
 
 tap.test(
@@ -80,3 +76,9 @@ tap.test(
  t.equal(event.input, undefined, 'input should be empty')
  }
 )
+
+tap.test('should capture token_count when callback is defined', async (t) => {
+ t.context.agent.llm.tokenCountCallback = () => 3
+ const event = new LlmEmbedding(t.context)
+ t.equal(event.token_count, 3)
+})

tests/unit/llm/event.tap.js

@@ -40,7 +40,6 @@ tap.beforeEach((t) => {
  t.context.segment = {
  id: 'segment-1',
  transaction: {
- id: 'tx-1',
  traceId: 'trace-1'
  }
  }
@@ -62,7 +61,6 @@ tap.test('create creates a new instance', async (t) => {
  t.equal(event.ingest_source, 'Node')
  t.equal(event.appName, 'test-app')
  t.equal(event.span_id, 'segment-1')
- t.equal(event.transaction_id, 'tx-1')
  t.equal(event.trace_id, 'trace-1')
  t.equal(event.request_id, 'request-1')
  t.equal(event['response.model'], 'model-1')