From c2a0314f0c0090763fb6be77808a75e72df3a0a4 Mon Sep 17 00:00:00 2001 From: James Sumners Date: Thu, 15 Feb 2024 14:16:57 -0500 Subject: [PATCH] chore: Refactor LLM user added metadata (#256) --- lib/llm/chat-completion-message.js | 3 +- lib/llm/chat-completion-summary.js | 3 +- lib/llm/event.js | 18 +++++++---- lib/v3/bedrock.js | 32 ++++--------------- tests/unit/llm/chat-completion-message.tap.js | 8 ++--- tests/unit/llm/chat-completion-summary.tap.js | 12 +++---- tests/unit/llm/event.tap.js | 5 ++- tests/versioned/common.js | 2 +- 8 files changed, 34 insertions(+), 49 deletions(-) diff --git a/lib/llm/chat-completion-message.js b/lib/llm/chat-completion-message.js index 0881cce..1005084 100644 --- a/lib/llm/chat-completion-message.js +++ b/lib/llm/chat-completion-message.js @@ -39,10 +39,9 @@ class LlmChatCompletionMessage extends LlmEvent { params = Object.assign({}, defaultParams, params) super(params) - const { agent, content, isResponse, index, completionId } = params + const { content, isResponse, index, completionId } = params this.is_response = isResponse - this.conversation_id = this.conversationId(agent) this.completion_id = completionId this.sequence = index this.content = content diff --git a/lib/llm/chat-completion-summary.js b/lib/llm/chat-completion-summary.js index ac17fc0..8450055 100644 --- a/lib/llm/chat-completion-summary.js +++ b/lib/llm/chat-completion-summary.js @@ -23,9 +23,8 @@ class LlmChatCompletionSummary extends LlmEvent { constructor(params = defaultParams) { super(params) - const { agent, segment, isError } = params + const { segment, isError } = params this.error = isError - this.conversation_id = this.conversationId(agent) this.duration = segment.getDurationInMillis() this['request.max_tokens'] = this.bedrockCommand.maxTokens diff --git a/lib/llm/event.js b/lib/llm/event.js index ce6b62d..56030de 100644 --- a/lib/llm/event.js +++ b/lib/llm/event.js @@ -59,6 +59,7 @@ class LlmEvent { this.transaction_id = segment.transaction.id this.trace_id = segment.transaction.traceId this.request_id = this.bedrockResponse.requestId + this.metadata = agent this['response.model'] = this.bedrockCommand.modelId this['request.model'] = this.bedrockCommand.modelId @@ -66,18 +67,21 @@ class LlmEvent { } /** - * Retrieve the conversation identifier from the custom attributes - * stored in the current transaction. + * Pull user set `llm.*` attributes from the current transaction and + * add them to the event. * * @param {object} agent The New Relic agent that provides access to the * transaction. - * - * @returns {string} */ - conversationId(agent) { + set metadata(agent) { const tx = agent.tracer.getTransaction() - const attrs = tx?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE) - return attrs?.['llm.conversation_id'] + const attrs = tx?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE) || {} + for (const [k, v] of Object.entries(attrs)) { + if (k.startsWith('llm.') === false) { + continue + } + this[k] = v + } } /** diff --git a/lib/v3/bedrock.js b/lib/v3/bedrock.js index 05784a4..3393374 100644 --- a/lib/v3/bedrock.js +++ b/lib/v3/bedrock.js @@ -35,33 +35,13 @@ function shouldSkipInstrumentation(config) { * @param {Agent} params.agent NR agent instance * @param {string} params.type LLM event type * @param {object} params.msg LLM event - * @param {object} params.tx The current transaction. */ -function recordEvent({ agent, type, msg, tx }) { +function recordEvent({ agent, type, msg }) { agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount() - setMetadata({ tx, msg }) msg.serialize() agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } -/** - * Decorate the outgoing message with metadata that has been set by the - * customer via `api.addCustomAttribute()`. - * - * @param {object} params Function parameters. - * @param {object} params.tx The current transaction. - * @param {object} params.msg The outgoing message to decorate. - */ -function setMetadata({ tx, msg }) { - const CONVERSATION_ID = 'llm.conversation_id' - const attrs = tx.trace?.custom?.get(DESTINATIONS.TRANS_SCOPE) || {} - for (const [key, value] of Object.entries(attrs)) { - if (key.startsWith('llm.') && key !== CONVERSATION_ID) { - msg[key] = value - } - } -} - /** * Assigns requestId, conversationId and messageIds for a given * chat completion response on the active transaction. @@ -78,7 +58,7 @@ function assignIdsToTx({ tx, msg, responseId }) { tracker.get(responseId) ?? new LlmTrackedIds({ requestId: msg.request_id, - conversationId: msg.conversation_id + conversationId: msg['llm.conversation_id'] }) trackedIds.message_ids.push(msg.id) tracker.set(responseId, trackedIds) @@ -115,7 +95,7 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR completionId: summary.id }) assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg }) - recordEvent({ agent, type: 'LlmChatCompletionMessage', msg, tx }) + recordEvent({ agent, type: 'LlmChatCompletionMessage', msg }) bedrockResponse.completions.forEach((content, index) => { const chatCompletionMessage = new LlmChatCompletionMessage({ @@ -129,10 +109,10 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR completionId: summary.id }) assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg: chatCompletionMessage }) - recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage, tx }) + recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage }) }) - recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: summary, tx }) + recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: summary }) if (err) { const llmError = new LlmError({ bedrockResponse, err, summary }) @@ -159,7 +139,7 @@ function recordEmbeddingMessage({ agent, segment, bedrockCommand, bedrockRespons isError: err !== null }) - recordEvent({ agent, type: 'LlmEmbedding', msg: embedding, tx: agent.getTransaction() }) + recordEvent({ agent, type: 'LlmEmbedding', msg: embedding }) if (err) { const llmError = new LlmError({ bedrockResponse, err, embedding }) agent.errors.add(segment.transaction, err, llmError) diff --git a/tests/unit/llm/chat-completion-message.tap.js b/tests/unit/llm/chat-completion-message.tap.js index 3284dab..b11a403 100644 --- a/tests/unit/llm/chat-completion-message.tap.js +++ b/tests/unit/llm/chat-completion-message.tap.js @@ -75,7 +75,7 @@ tap.beforeEach((t) => { tap.test('create creates a non-response instance', async (t) => { const event = new LlmChatCompletionMessage(t.context) t.equal(event.is_response, false) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.completion_id, 'completion-1') t.equal(event.sequence, 0) t.equal(event.content, 'who are you') @@ -89,7 +89,7 @@ tap.test('create creates a titan response instance', async (t) => { t.context.isResponse = true const event = new LlmChatCompletionMessage(t.context) t.equal(event.is_response, true) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.completion_id, 'completion-1') t.equal(event.sequence, 0) t.equal(event.content, 'a response') @@ -104,7 +104,7 @@ tap.test('create creates a cohere response instance', async (t) => { t.context.bedrockResponse.id = 42 const event = new LlmChatCompletionMessage(t.context) t.equal(event.is_response, true) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.completion_id, 'completion-1') t.equal(event.sequence, 0) t.equal(event.content, 'a response') @@ -119,7 +119,7 @@ tap.test('create creates a ai21 response instance when response.id is undefined' delete t.context.bedrockResponse.id const event = new LlmChatCompletionMessage(t.context) t.equal(event.is_response, true) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.completion_id, 'completion-1') t.equal(event.sequence, 0) t.equal(event.content, 'a response') diff --git a/tests/unit/llm/chat-completion-summary.tap.js b/tests/unit/llm/chat-completion-summary.tap.js index b8900d1..e3ecbdd 100644 --- a/tests/unit/llm/chat-completion-summary.tap.js +++ b/tests/unit/llm/chat-completion-summary.tap.js @@ -80,7 +80,7 @@ tap.test('creates a basic summary', async (t) => { t.context.bedrockResponse.inputTokenCount = 0 t.context.bedrockResponse.outputTokenCount = 0 const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) @@ -94,7 +94,7 @@ tap.test('creates a basic summary', async (t) => { tap.test('creates an ai21 summary', async (t) => { t.context.bedrockCommand.isAi21 = () => true const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) @@ -108,7 +108,7 @@ tap.test('creates an ai21 summary', async (t) => { tap.test('creates an claude summary', async (t) => { t.context.bedrockCommand.isClaude = () => true const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) @@ -122,7 +122,7 @@ tap.test('creates an claude summary', async (t) => { tap.test('creates a cohere summary', async (t) => { t.context.bedrockCommand.isCohere = () => true const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) @@ -136,7 +136,7 @@ tap.test('creates a cohere summary', async (t) => { tap.test('creates a llama2 summary', async (t) => { t.context.bedrockCommand.isLlama2 = () => true const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) @@ -150,7 +150,7 @@ tap.test('creates a llama2 summary', async (t) => { tap.test('creates a titan summary', async (t) => { t.context.bedrockCommand.isTitan = () => true const event = new LlmChatCompletionSummary(t.context) - t.equal(event.conversation_id, 'conversation-1') + t.equal(event['llm.conversation_id'], 'conversation-1') t.equal(event.duration, 100) t.equal(event['request.max_tokens'], 25) t.equal(event['request.temperature'], 0.5) diff --git a/tests/unit/llm/event.tap.js b/tests/unit/llm/event.tap.js index 00e06b4..49a844a 100644 --- a/tests/unit/llm/event.tap.js +++ b/tests/unit/llm/event.tap.js @@ -26,7 +26,8 @@ tap.beforeEach((t) => { get(key) { t.equal(key, TRANS_SCOPE) return { - ['llm.conversation_id']: 'conversation-1' + ['llm.conversation_id']: 'conversation-1', + omit: 'me' } } } @@ -67,6 +68,8 @@ tap.test('create creates a new instance', async (t) => { t.equal(event['response.model'], 'model-1') t.equal(event['request.model'], 'model-1') t.equal(event['request.max_tokens'], null) + t.equal(event['llm.conversation_id'], 'conversation-1') + t.equal(event.omit, undefined) }) tap.test('serializes the event', (t) => { diff --git a/tests/versioned/common.js b/tests/versioned/common.js index 2dc991a..7a0d2b1 100644 --- a/tests/versioned/common.js +++ b/tests/versioned/common.js @@ -130,7 +130,7 @@ function assertChatCompletionSummary({ 'id': /[\w]{8}-[\w]{4}-[\w]{4}-[\w]{4}-[\w]{12}/, 'appName': 'New Relic for Node.js tests', 'request_id': 'eda0760a-c3f0-4fc1-9a1e-75559d642866', - 'conversation_id': 'convo-id', + 'llm.conversation_id': 'convo-id', 'trace_id': tx.traceId, 'span_id': tx.trace.root.children[0].id, 'transaction_id': tx.id,