diff --git a/lib/llm/bedrock-response.js b/lib/llm/bedrock-response.js index 4c32c09e25..90ca94907c 100644 --- a/lib/llm/bedrock-response.js +++ b/lib/llm/bedrock-response.js @@ -136,7 +136,7 @@ class BedrockResponse { * @returns {string} */ get requestId() { - return this.#innerOutput.requestId + return this.headers?.['x-amzn-requestid'] } /** diff --git a/lib/llm/event.js b/lib/llm/event.js index 1ef4ecea93..afb9385121 100644 --- a/lib/llm/event.js +++ b/lib/llm/event.js @@ -64,7 +64,7 @@ class LlmEvent { this.span_id = segment.id this.transaction_id = segment.transaction.id this.trace_id = segment.transaction.traceId - this.request_id = this.bedrockResponse.headers['x-amzn-requestid'] + this.request_id = this.bedrockResponse.requestId this['response.model'] = this.bedrockCommand.modelId this['request.model'] = this.bedrockCommand.modelId diff --git a/lib/llm/index.js b/lib/llm/index.js index 39f8c52a28..7dd577fae2 100644 --- a/lib/llm/index.js +++ b/lib/llm/index.js @@ -11,5 +11,6 @@ module.exports = { LlmChatCompletionMessage: require('./chat-completion-message'), LlmChatCompletionSummary: require('./chat-completion-summary'), LlmEmbedding: require('./embedding'), - LlmEvent: require('./event') + LlmEvent: require('./event'), + LlmTrackedIds: require('./tracked-ids') } diff --git a/lib/llm/tracked-ids.js b/lib/llm/tracked-ids.js new file mode 100644 index 0000000000..77496790f7 --- /dev/null +++ b/lib/llm/tracked-ids.js @@ -0,0 +1,27 @@ +/* + * Copyright 2024 New Relic Corporation. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +'use strict' + +/** + * Represents the identifiers created and collected when creating LLM + * chat completions. + * + * @property {string[]} [message_ids=[]] Internal identifiers for each message + * sent in the conversation. + * @property {string} [coversation_id=""] User defined identifier of the chat + * completion conversation. + * @property {string} request_id Identifier of the request from the remote + * service for a chat completion. Populated by the `x-request-id` header + * in the request. + * @public + */ +module.exports = class LlmTrackedIds { + constructor({ messageIds, conversationId = '', requestId } = {}) { + this.message_ids = messageIds ?? [] + this.conversation_id = conversationId + this.request_id = requestId + } +} diff --git a/lib/v3/bedrock.js b/lib/v3/bedrock.js index 7801a4cbde..e53d54af4a 100644 --- a/lib/v3/bedrock.js +++ b/lib/v3/bedrock.js @@ -8,6 +8,7 @@ const { LlmChatCompletionMessage, LlmChatCompletionSummary, LlmEmbedding, + LlmTrackedIds, BedrockCommand, BedrockResponse } = require('../llm') @@ -36,6 +37,28 @@ function recordEvent({ agent, type, msg }) { agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } +/** + * Assigns requestId, conversationId and messageIds for a given + * chat completion response on the active transaction. + * This is used for generating LlmFeedbackEvent via `api.recordLlmFeedbackEvent` + * + * @param {object} params input params + * @param {Transaction} params.tx active transaction + * @param {LlmChatCompletionMessage} params.completionMsg chat completion message + * @param {string} params.responseId id of response + */ +function assignIdsToTx({ tx, msg, responseId }) { + const tracker = tx.llm.responses + const trackedIds = + tracker.get(responseId) ?? + new LlmTrackedIds({ + requestId: msg.request_id, + conversationId: msg.conversation_id + }) + trackedIds.message_ids.push(msg.id) + tracker.set(responseId, trackedIds) +} + /** * Registers the specification for instrumentation bedrock calls * @@ -76,6 +99,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args) const bedrockResponse = new BedrockResponse({ bedrockCommand, response }) if (modelType === 'completion') { + const tx = segment.transaction const summary = new LlmChatCompletionSummary({ credentials, agent, @@ -92,6 +116,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args) index: 0, completionId: summary.id }) + assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg }) recordEvent({ agent, type: 'LlmChatCompletionMessage', msg }) bedrockResponse.completions.forEach((content, index) => { @@ -105,6 +130,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args) content, completionId: summary.id }) + assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg: chatCompletionMessage }) recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage }) }) diff --git a/tests/unit/llm/bedrock-response.tap.js b/tests/unit/llm/bedrock-response.tap.js index 973087f07e..f08e91c743 100644 --- a/tests/unit/llm/bedrock-response.tap.js +++ b/tests/unit/llm/bedrock-response.tap.js @@ -57,13 +57,13 @@ tap.beforeEach((t) => { response: { 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 } }, output: { - requestId: 'aws-request-1', body: new TextEncoder().encode('{"foo":"foo"}') } } @@ -100,7 +100,7 @@ tap.test('non-conforming response is handled gracefully', async (t) => { t.equal(res.id, undefined) t.equal(res.inputTokenCount, 0) t.equal(res.outputTokenCount, 0) - t.equal(res.requestId, 'aws-request-1') + t.equal(res.requestId, undefined) t.equal(res.statusCode, 200) }) diff --git a/tests/unit/llm/event.tap.js b/tests/unit/llm/event.tap.js index 7ac4f32562..c4a19d855e 100644 --- a/tests/unit/llm/event.tap.js +++ b/tests/unit/llm/event.tap.js @@ -46,9 +46,7 @@ tap.beforeEach((t) => { } t.context.bedrockResponse = { - headers: { - 'x-amzn-requestid': 'request-1' - } + requestId: 'request-1' } t.context.bedrockCommand = { diff --git a/tests/versioned/common.js b/tests/versioned/common.js index 3c3295662f..9a80bbe1fe 100644 --- a/tests/versioned/common.js +++ b/tests/versioned/common.js @@ -117,6 +117,7 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, tokenUsage, err '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', 'trace_id': tx.traceId, 'span_id': tx.trace.root.children[0].id, 'transaction_id': tx.id, @@ -128,6 +129,8 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, tokenUsage, err 'api_key_last_four_digits': 'E ID', 'response.number_of_messages': 2, 'response.choices.finish_reason': 'endoftext', + 'request.temperature': 0.5, + 'request.max_tokens': 100, 'error': error } diff --git a/tests/versioned/v3/bedrock-chat-completions.tap.js b/tests/versioned/v3/bedrock-chat-completions.tap.js index 3ad453fbe1..3f39e2cc89 100644 --- a/tests/versioned/v3/bedrock-chat-completions.tap.js +++ b/tests/versioned/v3/bedrock-chat-completions.tap.js @@ -13,23 +13,26 @@ const createAiResponseServer = require('../aws-server-stubs/ai-server') const { FAKE_CREDENTIALS } = require('../aws-server-stubs') const requests = { ai21: (prompt, modelId) => ({ - body: JSON.stringify({ prompt }), + body: JSON.stringify({ prompt, temperature: 0.5, maxTokens: 100 }), modelId }), amazon: (prompt, modelId) => ({ - body: JSON.stringify({ inputText: prompt }), + body: JSON.stringify({ + inputText: prompt, + textGenerationConfig: { temperature: 0.5, maxTokenCount: 100 } + }), modelId }), claude: (prompt, modelId) => ({ - body: JSON.stringify({ prompt }), + body: JSON.stringify({ prompt, temperature: 0.5, max_tokens_to_sample: 100 }), modelId }), cohere: (prompt, modelId) => ({ - body: JSON.stringify({ prompt }), + body: JSON.stringify({ prompt, temperature: 0.5, max_tokens: 100 }), modelId }), llama2: (prompt, modelId) => ({ - body: JSON.stringify({ prompt }), + body: JSON.stringify({ prompt, max_gen_length: 100, temperature: 0.5 }), modelId }) } @@ -116,7 +119,9 @@ tap.afterEach(async (t) => { const command = new bedrock.InvokeModelCommand(input) const { agent } = helper + const api = helper.getAgentApi() helper.runInTransaction(async (tx) => { + api.addCustomAttribute('llm.conversation_id', 'convo-id') await client.send(command) const events = agent.customEventAggregator.events.toArray() t.equal(events.length, 3) @@ -153,4 +158,54 @@ tap.afterEach(async (t) => { t.equal(error.message, expected.body.message) } }) + + tap.test('should store ids and record feedback message accordingly', (t) => { + const { bedrock, client, helper } = t.context + const conversationId = 'convo-id' + const prompt = `text ${resKey} ultimate question` + const input = requests[resKey](prompt, modelId) + const command = new bedrock.InvokeModelCommand(input) + + const { agent } = helper + const api = helper.getAgentApi() + helper.runInTransaction(async (tx) => { + api.addCustomAttribute('llm.conversation_id', conversationId) + const response = await client.send(command) + const responseId = response.$metadata.requestId + const events = agent.customEventAggregator.events.toArray() + const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') + const ids = api.getLlmMessageIds({ responseId }) + const messageIds = chatMsgs.map((msg) => msg[1].id) + t.equal(ids.request_id, responseId) + t.equal(ids.conversation_id, conversationId) + // message_ids order varies over test run, sort them to assure consistency + t.same(ids.message_ids.sort(), messageIds.sort()) + api.recordLlmFeedbackEvent({ + conversationId: ids.conversation_id, + requestId: ids.request_id, + messageId: ids.message_ids[0], + category: 'test-event', + rating: '5 star', + message: 'You are a mathematician.', + metadata: { foo: 'foo' } + }) + const recordedEvents = agent.customEventAggregator.getEvents() + const [[, feedback]] = recordedEvents.filter(([{ type }]) => type === 'LlmFeedbackMessage') + + t.match(feedback, { + id: /[\w\d]{32}/, + conversation_id: ids.conversation_id, + request_id: ids.request_id, + message_id: ids.message_ids[0], + category: 'test-event', + rating: '5 star', + message: 'You are a mathematician.', + ingest_source: 'Node', + foo: 'foo' + }) + + tx.end() + t.end() + }) + }) })