From 236124018d272af4306f1cbfc2f2578e9bced59f Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 16 Jul 2024 12:12:30 -0700 Subject: [PATCH 01/33] feat: Set LLM events custom attributes --- api.js | 45 +++++++++++++++++++++++ lib/instrumentation/aws-sdk/v3/bedrock.js | 1 + lib/instrumentation/langchain/common.js | 1 + lib/instrumentation/openai.js | 1 + 4 files changed, 48 insertions(+) diff --git a/api.js b/api.js index 5c0f05b6dd..cde25b12e6 100644 --- a/api.js +++ b/api.js @@ -70,6 +70,8 @@ const CUSTOM_EVENT_TYPE_REGEX = /^[a-zA-Z0-9:_ ]+$/ * @class */ function API(agent) { + // throw 'Test'; + console.log('debugging agent') this.agent = agent this.shim = new TransactionShim(agent, 'NewRelicAPI') this.awsLambda = new AwsLambda(agent) @@ -1902,4 +1904,47 @@ API.prototype.ignoreApdex = function ignoreApdex() { transaction.ignoreApdex = true } +/** + * @callback setLlmCustomAttributesCallback + * @param {string} type LLM event type + * @param {object} msg LLM event + * @returns {object} Returns key/value pairs of attributes to be added + */ + +/** + * Add custom attributes to the LLM event. + * + * See documentation for newrelic.setLlmCustomAttributes for more information. + * + * An example of setting a custom attribute: + * + * newrelic.setLlmCustomAttributes((type, llmEvent) => { + * llmEvent['llm.testAttribute'] = 'testValue' + * return llmEvent + * }) + * + * @param {setLlmCustomAttributesCallback} callback A callback to handle recording of every LLM event. + */ +API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) { + this.agent.on('recordLlmEvent', function handler(type, msg) { + // console.log('record llm events fired', type, msg) + const attributes = callback(type, msg) + + for (const key in attributes) { + if (Object.hasOwn(attributes, key)) { + const value = attributes[key] + if (typeof value === 'object' || typeof value === 'function') { + delete attributes[key] + } else if (key.indexOf('llm.') !== 0) { + delete attributes[key] + attributes[`llm.${key}`] = value + } + } + } + + Object.assign(msg, attributes || {}) + // console.log('updated event', msg); + }) +} + module.exports = API diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 055c572c63..0bd3ee9361 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -55,6 +55,7 @@ function isStreamingEnabled({ commandName, config }) { */ function recordEvent({ agent, type, msg }) { msg.serialize() + agent.emit('recordLlmEvent', type, msg) agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } diff --git a/lib/instrumentation/langchain/common.js b/lib/instrumentation/langchain/common.js index b8e5de272f..1a30ec011e 100644 --- a/lib/instrumentation/langchain/common.js +++ b/lib/instrumentation/langchain/common.js @@ -49,6 +49,7 @@ common.mergeMetadata = function mergeMetadata(localMeta = {}, paramsMeta = {}) { */ common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) { agent.metrics.getOrCreateMetric(`${LANGCHAIN.TRACKING_PREFIX}/${pkgVersion}`).incrementCallCount() + agent.emit('recordLlmEvent', type, msg) agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index da0b4a22f4..b1fb60c309 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -75,6 +75,7 @@ function decorateSegment({ shim, result, apiKey }) { * @param {object} params.msg LLM event */ function recordEvent({ agent, type, msg }) { + agent.emit('recordLlmEvent', type, msg) agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } From 4dd3be822318650fb9dc3b78fbd0c0d76f3c4467 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 16 Jul 2024 12:15:36 -0700 Subject: [PATCH 02/33] fix: Debug statements --- api.js | 2 -- 1 file changed, 2 deletions(-) diff --git a/api.js b/api.js index cde25b12e6..bd7a0e16bb 100644 --- a/api.js +++ b/api.js @@ -70,8 +70,6 @@ const CUSTOM_EVENT_TYPE_REGEX = /^[a-zA-Z0-9:_ ]+$/ * @class */ function API(agent) { - // throw 'Test'; - console.log('debugging agent') this.agent = agent this.shim = new TransactionShim(agent, 'NewRelicAPI') this.awsLambda = new AwsLambda(agent) From 55da778d3feea96b6313df535eba07ab36dc2de3 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 16 Jul 2024 12:30:58 -0700 Subject: [PATCH 03/33] fix: Example --- api.js | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/api.js b/api.js index bd7a0e16bb..dd8ce70662 100644 --- a/api.js +++ b/api.js @@ -1916,9 +1916,8 @@ API.prototype.ignoreApdex = function ignoreApdex() { * * An example of setting a custom attribute: * - * newrelic.setLlmCustomAttributes((type, llmEvent) => { - * llmEvent['llm.testAttribute'] = 'testValue' - * return llmEvent + * newrelic.setLlmCustomAttributes((type, msg) => { + * return {'llm.testAttribute': 'testValue'} * }) * * @param {setLlmCustomAttributesCallback} callback A callback to handle recording of every LLM event. From 720be07d8e74eb3431b0b56954fae85c1c8e5fd6 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 16 Jul 2024 14:16:20 -0700 Subject: [PATCH 04/33] fix: Add guards --- api.js | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/api.js b/api.js index dd8ce70662..5e7931e4cc 100644 --- a/api.js +++ b/api.js @@ -1924,15 +1924,25 @@ API.prototype.ignoreApdex = function ignoreApdex() { */ API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) { this.agent.on('recordLlmEvent', function handler(type, msg) { - // console.log('record llm events fired', type, msg) const attributes = callback(type, msg) + if (!attributes) { + return + } + if (typeof attributes !== 'object') { + logger.warn( + `Invalid attributes provided for ${type} event. Object with key/value pairs like {"llm.attributeName": "someValue"} is expected` + ) + return + } for (const key in attributes) { if (Object.hasOwn(attributes, key)) { const value = attributes[key] if (typeof value === 'object' || typeof value === 'function') { + logger.warn(`Invalid attribute type for ${key}. Skipped.`) delete attributes[key] } else if (key.indexOf('llm.') !== 0) { + logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`) delete attributes[key] attributes[`llm.${key}`] = value } @@ -1940,7 +1950,6 @@ API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) } Object.assign(msg, attributes || {}) - // console.log('updated event', msg); }) } From 878d6b5e55d81b0ca8a8ac726492a40061577db5 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 17 Jul 2024 12:19:17 -0700 Subject: [PATCH 05/33] fix: Add unit test --- test/unit/api/api-llm.test.js | 40 +++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index a2fb477d7c..7975973985 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,6 +121,46 @@ tap.test('Agent API LLM methods', (t) => { }) }) + t.test('setLlmCustom attributes', async (t) => { + const { api } = t.context + let [verifyType, verifyMsg] = [] + api.setLlmCustomAttributes((type, msg) => { + verifyType = type + verifyMsg = msg + }) + const rce = api.recordCustomEvent + let event + api.recordCustomEvent = (name, data) => { + event = { name, data } + return rce.call(api, name, data) + } + t.teardown(() => { + api.recordCustomEvent = rce + }) + + helper.runInTransaction(api.agent, () => { + const result = api.recordLlmFeedbackEvent({ + traceId: 'trace-id', + category: 'test-cat', + rating: '5 star', + metadata: { foo: 'foo' } + }) + t.equal(result, undefined) + t.equal(loggerMock.warn.callCount, 0) + t.equal(verifyType, 'LlmFeedbackMessage') + t.not.not.undefined(verifyMsg) + t.match(event.data, { + id: /[\w\d]{32}/, + trace_id: 'trace-id', + category: 'test-cat', + rating: '5 star', + message: '', + foo: 'foo', + ingest_source: 'Node' + }) + }) + }) + t.test('setLlmTokenCount should register callback to calculate token counts', async (t) => { const { api, agent } = t.context function callback(model, content) { From de8b498f9ece726498accf77f4ffc5b36c07655a Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Mon, 22 Jul 2024 15:34:54 -0700 Subject: [PATCH 06/33] feat: WithLlmCustomAttributes --- api.js | 13 +++++++++++++ lib/instrumentation/openai.js | 10 ++++++---- package.json | 1 + test/unit/api/api-llm.test.js | 27 +++++++++++++++++++-------- 4 files changed, 39 insertions(+), 12 deletions(-) diff --git a/api.js b/api.js index 5e7931e4cc..9191028df0 100644 --- a/api.js +++ b/api.js @@ -32,6 +32,7 @@ const obfuscate = require('./lib/util/sql/obfuscate') const { DESTINATIONS } = require('./lib/config/attribute-filter') const parse = require('module-details-from-path') const { isSimpleObject } = require('./lib/util/objects') +const { AsyncLocalStorage } = require('async_hooks') /* * @@ -1953,4 +1954,16 @@ API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) }) } +API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context, callback) { + const transaction = this.agent.tracer.getTransaction() + + if (!transaction) { + logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.') + return + } + + transaction._llmContext = new AsyncLocalStorage() + transaction._llmContext.run(context, callback) +} + module.exports = API diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index b1fb60c309..7a50b705a0 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -73,9 +73,11 @@ function decorateSegment({ shim, result, apiKey }) { * @param {Agent} params.agent NR agent instance * @param {string} params.type LLM event type * @param {object} params.msg LLM event + * @param params.segment */ -function recordEvent({ agent, type, msg }) { - agent.emit('recordLlmEvent', type, msg) +function recordEvent({ agent, segment, type, msg }) { + const context = segment.transaction._llmContext.getStore() + agent.emit('recordLlmEvent', type, msg, context) agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) } @@ -143,10 +145,10 @@ function recordChatCompletionMessages({ agent, shim, segment, request, response, message }) - recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: completionMsg }) + recordEvent({ agent, segment, type: 'LlmChatCompletionMessage', msg: completionMsg }) }) - recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: completionSummary }) + recordEvent({ agent, segment, type: 'LlmChatCompletionSummary', msg: completionSummary }) if (err) { const llmError = new LlmErrorMessage({ cause: err, summary: completionSummary, response }) diff --git a/package.json b/package.json index 3aee2b38a5..44b77b5cea 100644 --- a/package.json +++ b/package.json @@ -176,6 +176,7 @@ "test": "npm run integration && npm run unit", "third-party-updates": "oss third-party manifest --includeOptDeps && oss third-party notices --includeOptDeps && git add THIRD_PARTY_NOTICES.md third_party_manifest.json", "unit": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap --test-regex='(\\/|^test\\/unit\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", + "unit-llm": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap test/unit/api/api-llm.test.js --timeout=180 --no-coverage --reporter classic", "unit:scripts": "time c8 -o ./coverage/scripts-unit tap --test-regex='(\\/|^bin\\/test\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", "update-cross-agent-tests": "./bin/update-cats.sh", "versioned-tests": "./bin/run-versioned-tests.sh", diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index 7975973985..c58c41b3b2 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,35 +121,46 @@ tap.test('Agent API LLM methods', (t) => { }) }) + const wait = function (time) { + return new Promise((resolve) => { + setTimeout(() => { + resolve() + }, time) + }) + } + t.test('setLlmCustom attributes', async (t) => { const { api } = t.context let [verifyType, verifyMsg] = [] api.setLlmCustomAttributes((type, msg) => { + // console.log('here'); verifyType = type verifyMsg = msg }) - const rce = api.recordCustomEvent - let event - api.recordCustomEvent = (name, data) => { + // console.log(`---test---`); + // const rce = api.recordCustomEvent + // let event + /* api.recordCustomEvent = (name, data) => { event = { name, data } return rce.call(api, name, data) } t.teardown(() => { api.recordCustomEvent = rce - }) + })*/ - helper.runInTransaction(api.agent, () => { + helper.runInTransaction(api.agent, async () => { const result = api.recordLlmFeedbackEvent({ traceId: 'trace-id', category: 'test-cat', rating: '5 star', metadata: { foo: 'foo' } }) + await wait(10) t.equal(result, undefined) t.equal(loggerMock.warn.callCount, 0) t.equal(verifyType, 'LlmFeedbackMessage') - t.not.not.undefined(verifyMsg) - t.match(event.data, { + t.notSame(verifyMsg, undefined) + /* t.match(event.data, { id: /[\w\d]{32}/, trace_id: 'trace-id', category: 'test-cat', @@ -157,7 +168,7 @@ tap.test('Agent API LLM methods', (t) => { message: '', foo: 'foo', ingest_source: 'Node' - }) + })*/ }) }) From dc0261a3fc1963b541069b7db01d1964c99226ba Mon Sep 17 00:00:00 2001 From: Ryan Kadri Date: Tue, 23 Jul 2024 15:34:11 -0400 Subject: [PATCH 07/33] feat: Merge parent context into children --- api.js | 13 ++++++++----- lib/instrumentation/openai.js | 6 +++++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/api.js b/api.js index 9191028df0..50e312474c 100644 --- a/api.js +++ b/api.js @@ -1924,8 +1924,8 @@ API.prototype.ignoreApdex = function ignoreApdex() { * @param {setLlmCustomAttributesCallback} callback A callback to handle recording of every LLM event. */ API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) { - this.agent.on('recordLlmEvent', function handler(type, msg) { - const attributes = callback(type, msg) + this.agent.on('recordLlmEvent', function handler(type, msg, context) { + const attributes = callback(type, msg, context) if (!attributes) { return } @@ -1959,11 +1959,14 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context if (!transaction) { logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.') - return + return callback() } - transaction._llmContext = new AsyncLocalStorage() - transaction._llmContext.run(context, callback) + if (!transaction._llmContext) { + transaction._llmContext = new AsyncLocalStorage() + } + const fullContext = { ...(transaction._llmContext.getStore() ?? {}), ...context } + return transaction._llmContext.run(fullContext, callback) } module.exports = API diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index 7a50b705a0..dc19603515 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -78,7 +78,11 @@ function decorateSegment({ shim, result, apiKey }) { function recordEvent({ agent, segment, type, msg }) { const context = segment.transaction._llmContext.getStore() agent.emit('recordLlmEvent', type, msg, context) - agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) + + agent.customEventAggregator.add([ + { type, timestamp: Date.now() }, + { ...context, ...msg } + ]) } /** From 08927ec1cc01060ad4bef7262521408a349aabbb Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 31 Jul 2024 14:04:49 -0700 Subject: [PATCH 08/33] feat: Option 4 --- api.js | 82 ++++++++++++------------------- lib/instrumentation/openai.js | 20 +++++--- test/unit/api/api-llm.test.js | 90 ++++++++++++++++++----------------- test/unit/api/stub.test.js | 2 +- 4 files changed, 90 insertions(+), 104 deletions(-) diff --git a/api.js b/api.js index 50e312474c..9763981bf6 100644 --- a/api.js +++ b/api.js @@ -32,7 +32,6 @@ const obfuscate = require('./lib/util/sql/obfuscate') const { DESTINATIONS } = require('./lib/config/attribute-filter') const parse = require('module-details-from-path') const { isSimpleObject } = require('./lib/util/objects') -const { AsyncLocalStorage } = require('async_hooks') /* * @@ -1904,69 +1903,48 @@ API.prototype.ignoreApdex = function ignoreApdex() { } /** - * @callback setLlmCustomAttributesCallback - * @param {string} type LLM event type - * @param {object} msg LLM event - * @returns {object} Returns key/value pairs of attributes to be added - */ - -/** - * Add custom attributes to the LLM event. + * Runs a function synchronously within a provided LLM custom attributes context and returns its return value. * - * See documentation for newrelic.setLlmCustomAttributes for more information. + * See documentation for newrelic.withLlmCustomAttributes for more information. * * An example of setting a custom attribute: * - * newrelic.setLlmCustomAttributes((type, msg) => { - * return {'llm.testAttribute': 'testValue'} + * newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someVallue'}, () => { + * return; * }) - * - * @param {setLlmCustomAttributesCallback} callback A callback to handle recording of every LLM event. + * @param {Object} context LLM custom attributes context + * @param {Function} callback synchronous function called within the context */ -API.prototype.setLlmCustomAttributes = function setLlmCustomAttributes(callback) { - this.agent.on('recordLlmEvent', function handler(type, msg, context) { - const attributes = callback(type, msg, context) - if (!attributes) { - return - } - if (typeof attributes !== 'object') { - logger.warn( - `Invalid attributes provided for ${type} event. Object with key/value pairs like {"llm.attributeName": "someValue"} is expected` - ) - return - } - - for (const key in attributes) { - if (Object.hasOwn(attributes, key)) { - const value = attributes[key] - if (typeof value === 'object' || typeof value === 'function') { - logger.warn(`Invalid attribute type for ${key}. Skipped.`) - delete attributes[key] - } else if (key.indexOf('llm.') !== 0) { - logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`) - delete attributes[key] - attributes[`llm.${key}`] = value - } - } - } - - Object.assign(msg, attributes || {}) - }) -} - API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context, callback) { - const transaction = this.agent.tracer.getTransaction() + const metric = this.agent.metrics.getOrCreateMetric( + NAMES.SUPPORTABILITY.API + '/withLlmCustomAttributes' + ) + metric.incrementCallCount() + const transaction = this.agent.tracer.getTransaction() if (!transaction) { logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.') - return callback() + return callback?.() + } + + for (const key in context) { + if (Object.hasOwn(context, key)) { + const value = context[key] + if (typeof value === 'object' || typeof value === 'function') { + logger.warn(`Invalid attribute type for ${key}. Skipped.`) + delete context[key] + } else if (key.indexOf('llm.') !== 0) { + logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`) + delete context[key] + context[`llm.${key}`] = value + } + } } - if (!transaction._llmContext) { - transaction._llmContext = new AsyncLocalStorage() - } - const fullContext = { ...(transaction._llmContext.getStore() ?? {}), ...context } - return transaction._llmContext.run(fullContext, callback) + const parentContext = this.agent._contextManager.getContext() + + const fullContext = parentContext ? Object.assign(parentContext, context || {}) : context + return this.agent._contextManager.runInContext(fullContext, callback, this) } module.exports = API diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index dc19603515..ea0b71d931 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -73,15 +73,21 @@ function decorateSegment({ shim, result, apiKey }) { * @param {Agent} params.agent NR agent instance * @param {string} params.type LLM event type * @param {object} params.msg LLM event - * @param params.segment */ -function recordEvent({ agent, segment, type, msg }) { - const context = segment.transaction._llmContext.getStore() - agent.emit('recordLlmEvent', type, msg, context) +function recordEvent({ agent, type, msg }) { + const context = agent._contextManager.getContext() + const llmContext = Object.keys(context || {}).reduce((result, key) => { + if (key.indexOf('llm.') === 0) { + result[key] = context[key] + } + return result + }, {}) + + agent.emit('recordLlmEvent', type, msg, llmContext) agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - { ...context, ...msg } + Object.assign(msg, llmContext || {}) ]) } @@ -149,10 +155,10 @@ function recordChatCompletionMessages({ agent, shim, segment, request, response, message }) - recordEvent({ agent, segment, type: 'LlmChatCompletionMessage', msg: completionMsg }) + recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: completionMsg }) }) - recordEvent({ agent, segment, type: 'LlmChatCompletionSummary', msg: completionSummary }) + recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: completionSummary }) if (err) { const llmError = new LlmErrorMessage({ cause: err, summary: completionSummary, response }) diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index c58c41b3b2..588bebb372 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,54 +121,56 @@ tap.test('Agent API LLM methods', (t) => { }) }) - const wait = function (time) { - return new Promise((resolve) => { - setTimeout(() => { - resolve() - }, time) - }) - } - - t.test('setLlmCustom attributes', async (t) => { + t.test('mimic withLlmCustomAttributes', async (t) => { const { api } = t.context - let [verifyType, verifyMsg] = [] - api.setLlmCustomAttributes((type, msg) => { - // console.log('here'); - verifyType = type - verifyMsg = msg + const contextManager = api.agent._contextManager + await helper.runInTransaction(api.agent, async () => { + await contextManager.runInContext({ test: 1 }, async () => { + const context = contextManager.getContext() + t.equal(context.test, 1) + }) }) - // console.log(`---test---`); - // const rce = api.recordCustomEvent - // let event - /* api.recordCustomEvent = (name, data) => { - event = { name, data } - return rce.call(api, name, data) - } - t.teardown(() => { - api.recordCustomEvent = rce - })*/ + }) - helper.runInTransaction(api.agent, async () => { - const result = api.recordLlmFeedbackEvent({ - traceId: 'trace-id', - category: 'test-cat', - rating: '5 star', - metadata: { foo: 'foo' } + t.test('withLlmCustomAttributes attributes', (t) => { + const { api } = t.context + helper.runInTransaction(api.agent, (tx) => { + const contextManager = api.agent._contextManager + t.context.agent.tracer.getTransaction = () => { + return tx + } + contextManager.setContext({ + initial: true }) - await wait(10) - t.equal(result, undefined) - t.equal(loggerMock.warn.callCount, 0) - t.equal(verifyType, 'LlmFeedbackMessage') - t.notSame(verifyMsg, undefined) - /* t.match(event.data, { - id: /[\w\d]{32}/, - trace_id: 'trace-id', - category: 'test-cat', - rating: '5 star', - message: '', - foo: 'foo', - ingest_source: 'Node' - })*/ + + api.withLlmCustomAttributes( + { + 'toRename': 'value1', + 'llm.number': 1, + 'llm.boolean': true, + 'toDelete': () => {}, + 'toDelete2': {}, + 'toDelete3': [] + }, + () => { + const parentContext = contextManager.getContext() + t.ok(parentContext.initial) + t.equal(parentContext['llm.toRename'], 'value1') + t.notOk(parentContext.toDelete) + t.notOk(parentContext.toDelete2) + t.notOk(parentContext.toDelete3) + t.equal(parentContext['llm.number'], 1) + t.equal(parentContext['llm.boolean'], true) + + api.withLlmCustomAttributes({ 'llm.someAttribute': 'someValue' }, () => { + const context = contextManager.getContext() + t.ok(context.initial) + t.equal(context[`llm.toRename`], 'value1') + t.equal(context['llm.someAttribute'], 'someValue') + t.end() + }) + } + ) }) }) diff --git a/test/unit/api/stub.test.js b/test/unit/api/stub.test.js index 68401c64a1..f31f0e9d07 100644 --- a/test/unit/api/stub.test.js +++ b/test/unit/api/stub.test.js @@ -8,7 +8,7 @@ const tap = require('tap') const API = require('../../../stub_api') -const EXPECTED_API_COUNT = 36 +const EXPECTED_API_COUNT = 37 tap.test('Agent API - Stubbed Agent API', (t) => { t.autoend() From 8b0cd841b3b3f947bdfb3ed2137baa62ac3f7b22 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 31 Jul 2024 14:08:05 -0700 Subject: [PATCH 09/33] =?UTF-8?q?fix:=20Unnecessary=C2=A0test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/unit/api/api-llm.test.js | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index 588bebb372..fc5125bd4d 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,17 +121,6 @@ tap.test('Agent API LLM methods', (t) => { }) }) - t.test('mimic withLlmCustomAttributes', async (t) => { - const { api } = t.context - const contextManager = api.agent._contextManager - await helper.runInTransaction(api.agent, async () => { - await contextManager.runInContext({ test: 1 }, async () => { - const context = contextManager.getContext() - t.equal(context.test, 1) - }) - }) - }) - t.test('withLlmCustomAttributes attributes', (t) => { const { api } = t.context helper.runInTransaction(api.agent, (tx) => { From 7ab7bcd30bf6f5b246a489a972d942e2588e8d6f Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 31 Jul 2024 14:30:27 -0700 Subject: [PATCH 10/33] fix: Test name --- test/unit/api/api-llm.test.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index fc5125bd4d..99c5c988ca 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,7 +121,7 @@ tap.test('Agent API LLM methods', (t) => { }) }) - t.test('withLlmCustomAttributes attributes', (t) => { + t.test('withLlmCustomAttributes', (t) => { const { api } = t.context helper.runInTransaction(api.agent, (tx) => { const contextManager = api.agent._contextManager From 0b2f2d614ad0d44a43530b8a7af379f6f7509552 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Thu, 1 Aug 2024 13:16:33 -0700 Subject: [PATCH 11/33] fix: Integration tests --- lib/instrumentation/aws-sdk/v3/bedrock.js | 14 +++- lib/instrumentation/langchain/common.js | 14 +++- lib/instrumentation/openai.js | 2 - package.json | 3 + test/unit/instrumentation/openai.test.js | 22 ++++++ .../bedrock-chat-completions.tap.js | 41 ++++++----- test/versioned/langchain/runnables.tap.js | 69 ++++++++++--------- 7 files changed, 109 insertions(+), 56 deletions(-) diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 0bd3ee9361..8b0c9219ab 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -55,8 +55,18 @@ function isStreamingEnabled({ commandName, config }) { */ function recordEvent({ agent, type, msg }) { msg.serialize() - agent.emit('recordLlmEvent', type, msg) - agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) + const context = agent._contextManager.getContext() + const llmContext = Object.keys(context || {}).reduce((result, key) => { + if (key.indexOf('llm.') === 0) { + result[key] = context[key] + } + return result + }, {}) + + agent.customEventAggregator.add([ + { type, timestamp: Date.now() }, + Object.assign(msg, llmContext || {}) + ]) } /** diff --git a/lib/instrumentation/langchain/common.js b/lib/instrumentation/langchain/common.js index 1a30ec011e..e1bbc922d4 100644 --- a/lib/instrumentation/langchain/common.js +++ b/lib/instrumentation/langchain/common.js @@ -49,8 +49,18 @@ common.mergeMetadata = function mergeMetadata(localMeta = {}, paramsMeta = {}) { */ common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) { agent.metrics.getOrCreateMetric(`${LANGCHAIN.TRACKING_PREFIX}/${pkgVersion}`).incrementCallCount() - agent.emit('recordLlmEvent', type, msg) - agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg]) + const context = agent._contextManager.getContext() + const llmContext = Object.keys(context || {}).reduce((result, key) => { + if (key.indexOf('llm.') === 0) { + result[key] = context[key] + } + return result + }, {}) + + agent.customEventAggregator.add([ + { type, timestamp: Date.now() }, + Object.assign(msg, llmContext || {}) + ]) } /** diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index ea0b71d931..8baa079ab5 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -83,8 +83,6 @@ function recordEvent({ agent, type, msg }) { return result }, {}) - agent.emit('recordLlmEvent', type, msg, llmContext) - agent.customEventAggregator.add([ { type, timestamp: Date.now() }, Object.assign(msg, llmContext || {}) diff --git a/package.json b/package.json index 883993b4f5..c5d5eb894c 100644 --- a/package.json +++ b/package.json @@ -177,9 +177,12 @@ "third-party-updates": "oss third-party manifest --includeOptDeps && oss third-party notices --includeOptDeps && git add THIRD_PARTY_NOTICES.md third_party_manifest.json", "unit": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap --test-regex='(\\/|^test\\/unit\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", "unit-llm": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap test/unit/api/api-llm.test.js --timeout=180 --no-coverage --reporter classic", + "unit-openai": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap test/unit/instrumentation/openai.test.js --timeout=180 --no-coverage --reporter classic", + "unit:scripts": "time c8 -o ./coverage/scripts-unit tap --test-regex='(\\/|^bin\\/test\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", "update-cross-agent-tests": "./bin/update-cats.sh", "versioned-tests": "./bin/run-versioned-tests.sh", + "versioned-tests-selected": "./bin/run-versioned-tests.sh langchain", "update-changelog-version": "node ./bin/update-changelog-version", "checkout-external-versioned": "node ./test/versioned-external/checkout-external-tests.js", "versioned:internal:major": "VERSIONED_MODE=--major npm run versioned:internal", diff --git a/test/unit/instrumentation/openai.test.js b/test/unit/instrumentation/openai.test.js index 95be982840..c2eab2d774 100644 --- a/test/unit/instrumentation/openai.test.js +++ b/test/unit/instrumentation/openai.test.js @@ -119,5 +119,27 @@ test('openai unit tests', (t) => { t.equal(isWrapped, false, 'should not wrap chat completions create') t.end() }) + + t.test('should record LLM custom events with attributes', (t) => { + const { shim, agent, initialize } = t.context + shim.pkgVersion = '4.12.2' + const MockOpenAi = getMockModule() + agent.config.ai_monitoring.record_content = { enabled: true } + initialize(agent, MockOpenAi, 'openai', shim) + const completions = new MockOpenAi.Chat.Completions() + agent._contextManager.setContext({ initial: true }) + const api = helper.getAgentApi() + helper.runInTransaction(agent, () => { + api.withLlmCustomAttributes({ 'llm.attribute': `someValue` }, async () => { + await completions.create({ stream: false, messages: [{ role: 'user', content: 'Hello' }] }) + const events = agent.customEventAggregator.events.toArray() + const [[, message]] = events + t.notOk(message.initial) + t.equal(message['llm.attribute'], 'someValue') + t.end() + }) + }) + }) + t.end() }) diff --git a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js index 61b2bc9b42..131c290d96 100644 --- a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js +++ b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js @@ -135,25 +135,30 @@ tap.afterEach(async (t) => { const api = helper.getAgentApi() helper.runInTransaction(agent, async (tx) => { api.addCustomAttribute('llm.conversation_id', 'convo-id') - await client.send(command) - const events = agent.customEventAggregator.events.toArray() - t.equal(events.length, 3) - const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] - const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') - - t.llmMessages({ - modelId, - prompt, - resContent: '42', - tx, - expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null, - chatMsgs + api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { + await client.send(command) + const events = agent.customEventAggregator.events.toArray() + t.equal(events.length, 3) + const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] + const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') + + t.llmMessages({ + modelId, + prompt, + resContent: '42', + tx, + expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null, + chatMsgs + }) + + t.llmSummary({ tx, modelId, chatSummary }) + + const [, message] = chatSummary + t.ok(message['llm.contextAttribute']) + + tx.end() + t.end() }) - - t.llmSummary({ tx, modelId, chatSummary }) - - tx.end() - t.end() }) } ) diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index e6ced1cc4c..5d964db70d 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -52,7 +52,6 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { t.test('should create langchain events for every invoke call', (t) => { const { agent, prompt, outputParser, model } = t.context - helper.runInTransaction(agent, async (tx) => { const input = { topic: 'scientist' } const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } @@ -99,39 +98,45 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { 'should create langchain events for every invoke call on chat prompt + model + parser', (t) => { const { agent, prompt, outputParser, model } = t.context - + const api = helper.getAgentApi() + agent._contextManager.setContext({ initial: true }) helper.runInTransaction(agent, async (tx) => { - const input = { topic: 'scientist' } - const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } - - const chain = prompt.pipe(model).pipe(outputParser) - await chain.invoke(input, options) - - const events = agent.customEventAggregator.events.toArray() - - const langchainEvents = filterLangchainEvents(events) - const langChainMessageEvents = filterLangchainEventsByType( - langchainEvents, - 'LlmChatCompletionMessage' - ) - const langChainSummaryEvents = filterLangchainEventsByType( - langchainEvents, - 'LlmChatCompletionSummary' - ) - - t.langchainSummary({ - tx, - chatSummary: langChainSummaryEvents[0] - }) - - t.langchainMessages({ - tx, - chatMsgs: langChainMessageEvents, - chatSummary: langChainSummaryEvents[0][1] + api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { + const input = { topic: 'scientist' } + const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } + + const chain = prompt.pipe(model).pipe(outputParser) + await chain.invoke(input, options) + + const events = agent.customEventAggregator.events.toArray() + + const langchainEvents = filterLangchainEvents(events) + const langChainMessageEvents = filterLangchainEventsByType( + langchainEvents, + 'LlmChatCompletionMessage' + ) + const langChainSummaryEvents = filterLangchainEventsByType( + langchainEvents, + 'LlmChatCompletionSummary' + ) + + t.langchainSummary({ + tx, + chatSummary: langChainSummaryEvents[0] + }) + const [[, message]] = events + + t.ok(message['llm.contextAttribute']) + + t.langchainMessages({ + tx, + chatMsgs: langChainMessageEvents, + chatSummary: langChainSummaryEvents[0][1] + }) + + tx.end() + t.end() }) - - tx.end() - t.end() }) } ) From f75dabfc0e21f31047732a950838c6ad701eb245 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Fri, 2 Aug 2024 08:55:42 -0700 Subject: [PATCH 12/33] fix: Remove extra npm scripts --- package.json | 3 --- 1 file changed, 3 deletions(-) diff --git a/package.json b/package.json index c5d5eb894c..78ff218727 100644 --- a/package.json +++ b/package.json @@ -176,9 +176,6 @@ "test": "npm run integration && npm run unit", "third-party-updates": "oss third-party manifest --includeOptDeps && oss third-party notices --includeOptDeps && git add THIRD_PARTY_NOTICES.md third_party_manifest.json", "unit": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap --test-regex='(\\/|^test\\/unit\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", - "unit-llm": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap test/unit/api/api-llm.test.js --timeout=180 --no-coverage --reporter classic", - "unit-openai": "rm -f newrelic_agent.log && time c8 -o ./coverage/unit tap test/unit/instrumentation/openai.test.js --timeout=180 --no-coverage --reporter classic", - "unit:scripts": "time c8 -o ./coverage/scripts-unit tap --test-regex='(\\/|^bin\\/test\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", "update-cross-agent-tests": "./bin/update-cats.sh", "versioned-tests": "./bin/run-versioned-tests.sh", From 19a9e80b75ed357e7571e80e160f476d952a4e09 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Fri, 2 Aug 2024 09:13:07 -0700 Subject: [PATCH 13/33] fix: Remove extra npm scripts --- package.json | 1 - 1 file changed, 1 deletion(-) diff --git a/package.json b/package.json index 78ff218727..b7dfd79105 100644 --- a/package.json +++ b/package.json @@ -179,7 +179,6 @@ "unit:scripts": "time c8 -o ./coverage/scripts-unit tap --test-regex='(\\/|^bin\\/test\\/.*\\.test\\.js)$' --timeout=180 --no-coverage --reporter classic", "update-cross-agent-tests": "./bin/update-cats.sh", "versioned-tests": "./bin/run-versioned-tests.sh", - "versioned-tests-selected": "./bin/run-versioned-tests.sh langchain", "update-changelog-version": "node ./bin/update-changelog-version", "checkout-external-versioned": "node ./test/versioned-external/checkout-external-tests.js", "versioned:internal:major": "VERSIONED_MODE=--major npm run versioned:internal", From 24a13bd3ca952c2760e3d628c2a18a98794c5b85 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Fri, 2 Aug 2024 09:37:43 -0700 Subject: [PATCH 14/33] fix: PR feedback --- lib/instrumentation/aws-sdk/v3/bedrock.js | 2 +- lib/instrumentation/langchain/common.js | 2 +- lib/instrumentation/openai.js | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 8b0c9219ab..168194853d 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -65,7 +65,7 @@ function recordEvent({ agent, type, msg }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign(msg, llmContext || {}) + Object.assign({}, msg, llmContext || {}) ]) } diff --git a/lib/instrumentation/langchain/common.js b/lib/instrumentation/langchain/common.js index e1bbc922d4..ba6d4aa40b 100644 --- a/lib/instrumentation/langchain/common.js +++ b/lib/instrumentation/langchain/common.js @@ -59,7 +59,7 @@ common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign(msg, llmContext || {}) + Object.assign({}, msg, llmContext || {}) ]) } diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index 8baa079ab5..0ca5ff5118 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -85,7 +85,7 @@ function recordEvent({ agent, type, msg }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign(msg, llmContext || {}) + Object.assign({}, msg, llmContext || {}) ]) } From af4aafa1e1f26516294e5f8074192a714e61ecca Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Fri, 2 Aug 2024 10:23:39 -0700 Subject: [PATCH 15/33] fix: PR feedback --- api.js | 2 -- 1 file changed, 2 deletions(-) diff --git a/api.js b/api.js index 9763981bf6..f10ee17588 100644 --- a/api.js +++ b/api.js @@ -1905,8 +1905,6 @@ API.prototype.ignoreApdex = function ignoreApdex() { /** * Runs a function synchronously within a provided LLM custom attributes context and returns its return value. * - * See documentation for newrelic.withLlmCustomAttributes for more information. - * * An example of setting a custom attribute: * * newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someVallue'}, () => { From d428aaa3abb0f8b6ac8493b64aebdf3c9c9cc031 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Mon, 5 Aug 2024 11:48:21 -0700 Subject: [PATCH 16/33] fix: PR feedback --- test/versioned/langchain/runnables.tap.js | 88 +++++++++++++---------- 1 file changed, 52 insertions(+), 36 deletions(-) diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index 5d964db70d..81b910ea71 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -94,49 +94,65 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { }) }) + t.test('should support custom attributes on the LLM events', (t) => { + const { agent, prompt, outputParser, model } = t.context + const api = helper.getAgentApi() + helper.runInTransaction(agent, async (tx) => { + api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { + const input = { topic: 'scientist' } + const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } + + const chain = prompt.pipe(model).pipe(outputParser) + await chain.invoke(input, options) + + const events = agent.customEventAggregator.events.toArray() + const [[, message]] = events + + t.ok(message['llm.contextAttribute']) + + tx.end() + t.end() + }) + }) + }) + t.test( 'should create langchain events for every invoke call on chat prompt + model + parser', (t) => { const { agent, prompt, outputParser, model } = t.context - const api = helper.getAgentApi() agent._contextManager.setContext({ initial: true }) helper.runInTransaction(agent, async (tx) => { - api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { - const input = { topic: 'scientist' } - const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } - - const chain = prompt.pipe(model).pipe(outputParser) - await chain.invoke(input, options) - - const events = agent.customEventAggregator.events.toArray() - - const langchainEvents = filterLangchainEvents(events) - const langChainMessageEvents = filterLangchainEventsByType( - langchainEvents, - 'LlmChatCompletionMessage' - ) - const langChainSummaryEvents = filterLangchainEventsByType( - langchainEvents, - 'LlmChatCompletionSummary' - ) - - t.langchainSummary({ - tx, - chatSummary: langChainSummaryEvents[0] - }) - const [[, message]] = events - - t.ok(message['llm.contextAttribute']) - - t.langchainMessages({ - tx, - chatMsgs: langChainMessageEvents, - chatSummary: langChainSummaryEvents[0][1] - }) - - tx.end() - t.end() + const input = { topic: 'scientist' } + const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } + + const chain = prompt.pipe(model).pipe(outputParser) + await chain.invoke(input, options) + + const events = agent.customEventAggregator.events.toArray() + + const langchainEvents = filterLangchainEvents(events) + const langChainMessageEvents = filterLangchainEventsByType( + langchainEvents, + 'LlmChatCompletionMessage' + ) + const langChainSummaryEvents = filterLangchainEventsByType( + langchainEvents, + 'LlmChatCompletionSummary' + ) + + t.langchainSummary({ + tx, + chatSummary: langChainSummaryEvents[0] }) + + t.langchainMessages({ + tx, + chatMsgs: langChainMessageEvents, + chatSummary: langChainSummaryEvents[0][1] + }) + + tx.end() + t.end() }) } ) From d2264ad2d46ceb8adbe8ad991409a26db30393c6 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Mon, 5 Aug 2024 15:03:01 -0700 Subject: [PATCH 17/33] fix: Unit test and pr feedback --- test/unit/instrumentation/openai.test.js | 21 ----- test/versioned/langchain/runnables.tap.js | 8 +- test/versioned/openai/chat-completions.tap.js | 83 +++++++++++++++++++ 3 files changed, 88 insertions(+), 24 deletions(-) diff --git a/test/unit/instrumentation/openai.test.js b/test/unit/instrumentation/openai.test.js index c2eab2d774..49a8c2aace 100644 --- a/test/unit/instrumentation/openai.test.js +++ b/test/unit/instrumentation/openai.test.js @@ -120,26 +120,5 @@ test('openai unit tests', (t) => { t.end() }) - t.test('should record LLM custom events with attributes', (t) => { - const { shim, agent, initialize } = t.context - shim.pkgVersion = '4.12.2' - const MockOpenAi = getMockModule() - agent.config.ai_monitoring.record_content = { enabled: true } - initialize(agent, MockOpenAi, 'openai', shim) - const completions = new MockOpenAi.Chat.Completions() - agent._contextManager.setContext({ initial: true }) - const api = helper.getAgentApi() - helper.runInTransaction(agent, () => { - api.withLlmCustomAttributes({ 'llm.attribute': `someValue` }, async () => { - await completions.create({ stream: false, messages: [{ role: 'user', content: 'Hello' }] }) - const events = agent.customEventAggregator.events.toArray() - const [[, message]] = events - t.notOk(message.initial) - t.equal(message['llm.attribute'], 'someValue') - t.end() - }) - }) - }) - t.end() }) diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index 81b910ea71..8daeada033 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -104,10 +104,12 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { const chain = prompt.pipe(model).pipe(outputParser) await chain.invoke(input, options) - const events = agent.customEventAggregator.events.toArray() - const [[, message]] = events - + const responses = events.filter((event) => { + const [, chainEvent] = event + return chainEvent.virtual_llm === true && chainEvent.is_response === true + }) + const [[, message]] = responses t.ok(message['llm.contextAttribute']) tx.end() diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js index 67334d7fd9..3af85e1e41 100644 --- a/test/versioned/openai/chat-completions.tap.js +++ b/test/versioned/openai/chat-completions.tap.js @@ -448,4 +448,87 @@ tap.test('OpenAI instrumentation - chat completions', (t) => { test.end() }) }) + + t.test('1should create chat completion message and summary for every message sent', (test) => { + const { client, agent } = t.context + helper.runInTransaction(agent, async (tx) => { + const model = 'gpt-3.5-turbo-0613' + const content = 'You are a mathematician.' + await client.chat.completions.create({ + max_tokens: 100, + temperature: 0.5, + model, + messages: [ + { role: 'user', content }, + { role: 'user', content: 'What does 1 plus 1 equal?' } + ] + }) + + const events = agent.customEventAggregator.events.toArray() + test.equal(events.length, 4, 'should create a chat completion message and summary event') + const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') + test.llmMessages({ + tx, + chatMsgs, + model, + id: 'chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat', + resContent: '1 plus 2 is 3.', + reqContent: content + }) + + const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] + test.llmSummary({ tx, model, chatSummary, tokenUsage: true }) + tx.end() + test.end() + }) + }) + + t.test('should record LLM custom events with attributes', (test) => { + const { client, agent } = t.context + const api = helper.getAgentApi() + + helper.runInTransaction(agent, () => { + api.withLlmCustomAttributes({ 'llm.initial': true }, async () => { + await api.withLlmCustomAttributes({ 'llm.attribute': `someValue` }, async () => { + agent.config.ai_monitoring.streaming.enabled = true + const model = 'gpt-3.5-turbo-0613' + const content = 'You are a mathematician.' + await client.chat.completions.create({ + max_tokens: 100, + temperature: 0.5, + model, + messages: [ + { role: 'user', content }, + { role: 'user', content: 'What does 1 plus 1 equal?' } + ] + }) + const events = agent.customEventAggregator.events.toArray() + const [[, message]] = events + t.ok(message['llm.initial']) + t.equal(message['llm.attribute'], 'someValue') + }) + + await api.withLlmCustomAttributes({ 'llm.anotherFlow': 'anotherValue' }, async () => { + agent.config.ai_monitoring.streaming.enabled = true + const model = 'gpt-3.5-turbo-0613' + const content = 'You are a mathematician.' + await client.chat.completions.create({ + max_tokens: 100, + temperature: 0.5, + model, + messages: [ + { role: 'user', content }, + { role: 'user', content: 'What does 1 plus 1 equal?' } + ] + }) + const events = agent.customEventAggregator.events.toArray() + const [[, message]] = events + t.ok(message['llm.initial']) + t.equal(message['llm.anotherFlow'], 'anotherValue') + }) + + test.end() + }) + }) + }) }) From ef488dab3b26667283b7ccba6da209675359b95f Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 13 Aug 2024 08:33:33 -0700 Subject: [PATCH 18/33] fix: PR feedback --- .../bedrock-chat-completions.tap.js | 60 +++++++++++++++++++ test/versioned/langchain/runnables.tap.js | 5 ++ test/versioned/openai/chat-completions.tap.js | 2 +- 3 files changed, 66 insertions(+), 1 deletion(-) diff --git a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js index 131c290d96..2f15100d38 100644 --- a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js +++ b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js @@ -163,6 +163,66 @@ tap.afterEach(async (t) => { } ) + tap.test( + `${modelId}: properly create the LlmChatCompletionMessage(s) and LlmChatCompletionSummary events`, + (t) => { + const { bedrock, client, agent } = t.context + const prompt = `text ${resKey} ultimate question` + const input = requests[resKey](prompt, modelId) + const command = new bedrock.InvokeModelCommand(input) + + const api = helper.getAgentApi() + helper.runInTransaction(agent, async (tx) => { + api.addCustomAttribute('llm.conversation_id', 'convo-id') + await client.send(command) + const events = agent.customEventAggregator.events.toArray() + t.equal(events.length, 3) + const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] + const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') + + t.llmMessages({ + modelId, + prompt, + resContent: '42', + tx, + expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null, + chatMsgs + }) + + t.llmSummary({ tx, modelId, chatSummary }) + + tx.end() + t.end() + }) + } + ) + + tap.test( + `${modelId}: supports custom attributes on LlmChatCompletionMessage(s) and LlmChatCompletionSummary events`, + (t) => { + const { bedrock, client, agent } = t.context + const prompt = `text ${resKey} ultimate question` + const input = requests[resKey](prompt, modelId) + const command = new bedrock.InvokeModelCommand(input) + + const api = helper.getAgentApi() + helper.runInTransaction(agent, async (tx) => { + api.addCustomAttribute('llm.conversation_id', 'convo-id') + api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { + await client.send(command) + const events = agent.customEventAggregator.events.toArray() + + const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] + const [, message] = chatSummary + t.ok(message['llm.contextAttribute']) + + tx.end() + t.end() + }) + }) + } + ) + tap.test(`${modelId}: text answer (streamed)`, (t) => { if (modelId.includes('ai21')) { t.skip('model does not support streaming') diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index 8daeada033..5b2e7dcaee 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -105,10 +105,15 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { const chain = prompt.pipe(model).pipe(outputParser) await chain.invoke(input, options) const events = agent.customEventAggregator.events.toArray() + + // Issue: Custom attributes present only for events with `virtual_llm` flag + const responses = events + /* // Workarround const responses = events.filter((event) => { const [, chainEvent] = event return chainEvent.virtual_llm === true && chainEvent.is_response === true }) + */ const [[, message]] = responses t.ok(message['llm.contextAttribute']) diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js index 3af85e1e41..ded31e4df1 100644 --- a/test/versioned/openai/chat-completions.tap.js +++ b/test/versioned/openai/chat-completions.tap.js @@ -449,7 +449,7 @@ tap.test('OpenAI instrumentation - chat completions', (t) => { }) }) - t.test('1should create chat completion message and summary for every message sent', (test) => { + t.test('should create chat completion message and summary for every message sent', (test) => { const { client, agent } = t.context helper.runInTransaction(agent, async (tx) => { const model = 'gpt-3.5-turbo-0613' From bcc9e15f50f07630bcebb7d3441195e1cdfe0844 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 13 Aug 2024 08:35:44 -0700 Subject: [PATCH 19/33] fix: Typo --- api.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api.js b/api.js index f10ee17588..448ffcc367 100644 --- a/api.js +++ b/api.js @@ -1907,7 +1907,7 @@ API.prototype.ignoreApdex = function ignoreApdex() { * * An example of setting a custom attribute: * - * newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someVallue'}, () => { + * newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someValue'}, () => { * return; * }) * @param {Object} context LLM custom attributes context From cef42898b5abdaa5eb2b64f70a730be0d3f99db3 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 13 Aug 2024 08:40:35 -0700 Subject: [PATCH 20/33] fix: PR feedback --- test/versioned/langchain/runnables.tap.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index 5b2e7dcaee..fefc01579f 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -115,7 +115,7 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { }) */ const [[, message]] = responses - t.ok(message['llm.contextAttribute']) + t.equal(message['llm.contextAttribute'], 'someValue') tx.end() t.end() From 87588e324c780dc50f9aee8ced29ca8c17b4712c Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 13 Aug 2024 08:42:36 -0700 Subject: [PATCH 21/33] fix: PR feedback --- .../bedrock-chat-completions.tap.js | 39 ------------------- 1 file changed, 39 deletions(-) diff --git a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js index 2f15100d38..1b63b3f659 100644 --- a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js +++ b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js @@ -124,45 +124,6 @@ tap.afterEach(async (t) => { }) }) - tap.test( - `${modelId}: properly create the LlmChatCompletionMessage(s) and LlmChatCompletionSummary events`, - (t) => { - const { bedrock, client, agent } = t.context - const prompt = `text ${resKey} ultimate question` - const input = requests[resKey](prompt, modelId) - const command = new bedrock.InvokeModelCommand(input) - - const api = helper.getAgentApi() - helper.runInTransaction(agent, async (tx) => { - api.addCustomAttribute('llm.conversation_id', 'convo-id') - api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => { - await client.send(command) - const events = agent.customEventAggregator.events.toArray() - t.equal(events.length, 3) - const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] - const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') - - t.llmMessages({ - modelId, - prompt, - resContent: '42', - tx, - expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null, - chatMsgs - }) - - t.llmSummary({ tx, modelId, chatSummary }) - - const [, message] = chatSummary - t.ok(message['llm.contextAttribute']) - - tx.end() - t.end() - }) - }) - } - ) - tap.test( `${modelId}: properly create the LlmChatCompletionMessage(s) and LlmChatCompletionSummary events`, (t) => { From d0101f685ca13b757ce2ade7f5e9831c3928c3dc Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 13 Aug 2024 08:46:23 -0700 Subject: [PATCH 22/33] fix: Unit test --- test/versioned/langchain/runnables.tap.js | 1 - 1 file changed, 1 deletion(-) diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index fefc01579f..5cd15a2c55 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -127,7 +127,6 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { 'should create langchain events for every invoke call on chat prompt + model + parser', (t) => { const { agent, prompt, outputParser, model } = t.context - agent._contextManager.setContext({ initial: true }) helper.runInTransaction(agent, async (tx) => { const input = { topic: 'scientist' } const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] } From e7531d75dc2516f0b876bd8ea7f0575f0765771e Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 14 Aug 2024 12:51:26 -0700 Subject: [PATCH 23/33] fix: Apply solution 1 --- api.js | 8 +- lib/instrumentation/aws-sdk/v3/bedrock.js | 9 +- lib/instrumentation/langchain/common.js | 9 +- lib/instrumentation/openai.js | 9 +- lib/util/llm-utils.js | 42 ++++++++++ test/unit/api/api-llm.test.js | 12 +-- test/unit/util/llm-utils.test.js | 44 ++++++++++ .../bedrock-chat-completions.tap.js | 2 +- test/versioned/langchain/runnables.tap.js | 10 +-- test/versioned/openai/chat-completions.tap.js | 83 +++++++++++-------- 10 files changed, 150 insertions(+), 78 deletions(-) create mode 100644 lib/util/llm-utils.js create mode 100644 test/unit/util/llm-utils.test.js diff --git a/api.js b/api.js index 448ffcc367..336018b556 100644 --- a/api.js +++ b/api.js @@ -32,6 +32,7 @@ const obfuscate = require('./lib/util/sql/obfuscate') const { DESTINATIONS } = require('./lib/config/attribute-filter') const parse = require('module-details-from-path') const { isSimpleObject } = require('./lib/util/objects') +const { AsyncLocalStorage } = require('async_hooks') /* * @@ -1939,10 +1940,11 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context } } - const parentContext = this.agent._contextManager.getContext() + transaction._llmContextManager = transaction._llmContextManager || new AsyncLocalStorage() + const parentContext = transaction._llmContextManager.getStore() - const fullContext = parentContext ? Object.assign(parentContext, context || {}) : context - return this.agent._contextManager.runInContext(fullContext, callback, this) + const fullContext = Object.assign({}, parentContext || {}, context || {}) + return transaction._llmContextManager.run(fullContext, callback) } module.exports = API diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 168194853d..2d6a6d85ee 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -18,6 +18,7 @@ const { DESTINATIONS } = require('../../../config/attribute-filter') const { AI } = require('../../../metrics/names') const { RecorderSpec } = require('../../../shim/specs') const InstrumentationDescriptor = require('../../../instrumentation-descriptor') +const { extractLlmContext } = require('../../../util/llm-utils') let TRACKING_METRIC @@ -55,13 +56,7 @@ function isStreamingEnabled({ commandName, config }) { */ function recordEvent({ agent, type, msg }) { msg.serialize() - const context = agent._contextManager.getContext() - const llmContext = Object.keys(context || {}).reduce((result, key) => { - if (key.indexOf('llm.') === 0) { - result[key] = context[key] - } - return result - }, {}) + const llmContext = extractLlmContext(agent) agent.customEventAggregator.add([ { type, timestamp: Date.now() }, diff --git a/lib/instrumentation/langchain/common.js b/lib/instrumentation/langchain/common.js index ba6d4aa40b..3644f77485 100644 --- a/lib/instrumentation/langchain/common.js +++ b/lib/instrumentation/langchain/common.js @@ -7,6 +7,7 @@ const { AI: { LANGCHAIN } } = require('../../metrics/names') +const { extractLlmContext } = require('../../util/llm-utils') const common = module.exports @@ -49,13 +50,7 @@ common.mergeMetadata = function mergeMetadata(localMeta = {}, paramsMeta = {}) { */ common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) { agent.metrics.getOrCreateMetric(`${LANGCHAIN.TRACKING_PREFIX}/${pkgVersion}`).incrementCallCount() - const context = agent._contextManager.getContext() - const llmContext = Object.keys(context || {}).reduce((result, key) => { - if (key.indexOf('llm.') === 0) { - result[key] = context[key] - } - return result - }, {}) + const llmContext = extractLlmContext(agent) agent.customEventAggregator.add([ { type, timestamp: Date.now() }, diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index 0ca5ff5118..c85c9758e9 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -12,6 +12,7 @@ const { LlmErrorMessage } = require('../../lib/llm-events/openai') const { RecorderSpec } = require('../../lib/shim/specs') +const { extractLlmContext } = require('../util/llm-utils') const MIN_VERSION = '4.0.0' const MIN_STREAM_VERSION = '4.12.2' @@ -75,13 +76,7 @@ function decorateSegment({ shim, result, apiKey }) { * @param {object} params.msg LLM event */ function recordEvent({ agent, type, msg }) { - const context = agent._contextManager.getContext() - const llmContext = Object.keys(context || {}).reduce((result, key) => { - if (key.indexOf('llm.') === 0) { - result[key] = context[key] - } - return result - }, {}) + const llmContext = extractLlmContext(agent) agent.customEventAggregator.add([ { type, timestamp: Date.now() }, diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js new file mode 100644 index 0000000000..d391b6d7ba --- /dev/null +++ b/lib/util/llm-utils.js @@ -0,0 +1,42 @@ +/* + * Copyright 2020 New Relic Corporation. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +'use strict' + +exports = module.exports = { extractLlmContext, extractLlmAttribtues } + +/** + * Extract LLM attributes from the LLM context + * + * @param {Object} context LLM context object + * @returns {Object} LLM custom attributes + */ +function extractLlmAttribtues(context) { + return Object.keys(context || {}).reduce((result, key) => { + if (key.indexOf('llm.') === 0) { + result[key] = context[key] + } + return result + }, {}) +} + +/** + * Extract LLM context from the active transaction + * + * @param {Agent} params.agent NR agent instance + * @param agent + * @returns {Object} LLM context object + */ +function extractLlmContext(agent) { + let context = agent.tracer.getTransaction()._llmContextManager + ? agent.tracer.getTransaction()._llmContextManager.getStore() + : null + const llmContext = {} + while (context) { + Object.assign(llmContext, extractLlmAttribtues(context)) + context = null + } + return llmContext +} diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index 99c5c988ca..772253fc2d 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -124,13 +124,9 @@ tap.test('Agent API LLM methods', (t) => { t.test('withLlmCustomAttributes', (t) => { const { api } = t.context helper.runInTransaction(api.agent, (tx) => { - const contextManager = api.agent._contextManager t.context.agent.tracer.getTransaction = () => { return tx } - contextManager.setContext({ - initial: true - }) api.withLlmCustomAttributes( { @@ -142,8 +138,8 @@ tap.test('Agent API LLM methods', (t) => { 'toDelete3': [] }, () => { - const parentContext = contextManager.getContext() - t.ok(parentContext.initial) + const contextManager = tx._llmContextManager + const parentContext = contextManager.getStore() t.equal(parentContext['llm.toRename'], 'value1') t.notOk(parentContext.toDelete) t.notOk(parentContext.toDelete2) @@ -152,8 +148,8 @@ tap.test('Agent API LLM methods', (t) => { t.equal(parentContext['llm.boolean'], true) api.withLlmCustomAttributes({ 'llm.someAttribute': 'someValue' }, () => { - const context = contextManager.getContext() - t.ok(context.initial) + const contextManager = tx._llmContextManager + const context = contextManager.getStore() t.equal(context[`llm.toRename`], 'value1') t.equal(context['llm.someAttribute'], 'someValue') t.end() diff --git a/test/unit/util/llm-utils.test.js b/test/unit/util/llm-utils.test.js new file mode 100644 index 0000000000..1e0ccc1c30 --- /dev/null +++ b/test/unit/util/llm-utils.test.js @@ -0,0 +1,44 @@ +/* + * Copyright 2023 New Relic Corporation. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +'use strict' + +const tap = require('tap') +const { extractLlmAttribtues, extractLlmContext } = require('../../../lib/util/llm-utils') +const { AsyncLocalStorage } = require('async_hooks') + +tap.test('extractLlmAttributes', (t) => { + const context = { + 'skip': 1, + 'llm.get': 2, + 'fllm.skip': 3 + } + + const llmContext = extractLlmAttribtues(context) + t.notOk(llmContext.skip) + t.notOk(llmContext['fllm.skip']) + t.equal(llmContext['llm.get'], 2) + t.end() +}) + +tap.test('extractLlmContext', (t) => { + const tx = { + _llmContextManager: new AsyncLocalStorage() + } + const agent = { + tracer: { + getTransaction: () => { + return tx + } + } + } + + tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => { + const llmContext = extractLlmContext(agent) + t.equal(llmContext['llm.test'], 1) + t.notOk(llmContext.skip) + t.end() + }) +}) diff --git a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js index 1b63b3f659..cda36a41cc 100644 --- a/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js +++ b/test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js @@ -175,7 +175,7 @@ tap.afterEach(async (t) => { const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] const [, message] = chatSummary - t.ok(message['llm.contextAttribute']) + t.equal(message['llm.contextAttribute'], 'someValue') tx.end() t.end() diff --git a/test/versioned/langchain/runnables.tap.js b/test/versioned/langchain/runnables.tap.js index 5cd15a2c55..a2d3aacc7f 100644 --- a/test/versioned/langchain/runnables.tap.js +++ b/test/versioned/langchain/runnables.tap.js @@ -106,15 +106,7 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => { await chain.invoke(input, options) const events = agent.customEventAggregator.events.toArray() - // Issue: Custom attributes present only for events with `virtual_llm` flag - const responses = events - /* // Workarround - const responses = events.filter((event) => { - const [, chainEvent] = event - return chainEvent.virtual_llm === true && chainEvent.is_response === true - }) - */ - const [[, message]] = responses + const [[, message]] = events t.equal(message['llm.contextAttribute'], 'someValue') tx.end() diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js index ded31e4df1..054a4a7231 100644 --- a/test/versioned/openai/chat-completions.tap.js +++ b/test/versioned/openai/chat-completions.tap.js @@ -488,43 +488,54 @@ tap.test('OpenAI instrumentation - chat completions', (t) => { const api = helper.getAgentApi() helper.runInTransaction(agent, () => { - api.withLlmCustomAttributes({ 'llm.initial': true }, async () => { - await api.withLlmCustomAttributes({ 'llm.attribute': `someValue` }, async () => { - agent.config.ai_monitoring.streaming.enabled = true - const model = 'gpt-3.5-turbo-0613' - const content = 'You are a mathematician.' - await client.chat.completions.create({ - max_tokens: 100, - temperature: 0.5, - model, - messages: [ - { role: 'user', content }, - { role: 'user', content: 'What does 1 plus 1 equal?' } - ] - }) - const events = agent.customEventAggregator.events.toArray() - const [[, message]] = events - t.ok(message['llm.initial']) - t.equal(message['llm.attribute'], 'someValue') - }) + api.withLlmCustomAttributes({ 'llm.shared': true, 'llm.path': 'root/' }, async () => { + await api.withLlmCustomAttributes( + { 'llm.path': 'root/branch1', 'llm.attr1': true }, + async () => { + agent.config.ai_monitoring.streaming.enabled = true + const model = 'gpt-3.5-turbo-0613' + const content = 'You are a mathematician.' + await client.chat.completions.create({ + max_tokens: 100, + temperature: 0.5, + model, + messages: [ + { role: 'user', content }, + { role: 'user', content: 'What does 1 plus 1 equal?' } + ] + }) + } + ) - await api.withLlmCustomAttributes({ 'llm.anotherFlow': 'anotherValue' }, async () => { - agent.config.ai_monitoring.streaming.enabled = true - const model = 'gpt-3.5-turbo-0613' - const content = 'You are a mathematician.' - await client.chat.completions.create({ - max_tokens: 100, - temperature: 0.5, - model, - messages: [ - { role: 'user', content }, - { role: 'user', content: 'What does 1 plus 1 equal?' } - ] - }) - const events = agent.customEventAggregator.events.toArray() - const [[, message]] = events - t.ok(message['llm.initial']) - t.equal(message['llm.anotherFlow'], 'anotherValue') + await api.withLlmCustomAttributes( + { 'llm.path': 'root/branch2', 'llm.attr2': true }, + async () => { + agent.config.ai_monitoring.streaming.enabled = true + const model = 'gpt-3.5-turbo-0613' + const content = 'You are a mathematician.' + await client.chat.completions.create({ + max_tokens: 100, + temperature: 0.5, + model, + messages: [ + { role: 'user', content }, + { role: 'user', content: 'What does 1 plus 2 equal?' } + ] + }) + } + ) + + const events = agent.customEventAggregator.events.toArray().map((event) => event[1]) + + events.forEach((event) => { + t.ok(event['llm.shared']) + if (event['llm.path'] === 'root/branch1') { + t.ok(event['llm.attr1']) + t.notOk(event['llm.attr2']) + } else { + t.ok(event['llm.attr2']) + t.notOk(event['llm.attr1']) + } }) test.end() From 308bdeb243683ec483c01f921e73721e52e3586b Mon Sep 17 00:00:00 2001 From: Webrealizer Date: Thu, 15 Aug 2024 08:01:07 -0700 Subject: [PATCH 24/33] Update lib/util/llm-utils.js PR feedback Co-authored-by: James Sumners --- lib/util/llm-utils.js | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js index d391b6d7ba..e73b9e20f6 100644 --- a/lib/util/llm-utils.js +++ b/lib/util/llm-utils.js @@ -25,8 +25,7 @@ function extractLlmAttribtues(context) { /** * Extract LLM context from the active transaction * - * @param {Agent} params.agent NR agent instance - * @param agent + * @param {Agent} agent NR agent instance * @returns {Object} LLM context object */ function extractLlmContext(agent) { From a32f932dacccab5137df9fa3c848769065959265 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Thu, 15 Aug 2024 09:01:02 -0700 Subject: [PATCH 25/33] fix: PR feedback --- api.js | 31 +++++++++++++++++-------------- lib/util/llm-utils.js | 10 +++------- test/unit/api/api-llm.test.js | 5 +++++ 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/api.js b/api.js index 336018b556..d7bf0775b0 100644 --- a/api.js +++ b/api.js @@ -1921,22 +1921,25 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context metric.incrementCallCount() const transaction = this.agent.tracer.getTransaction() + + if (!callback || typeof callback !== 'function') { + logger.warn('withLlmCustomAttributes must be used with a valid callback') + return + } + if (!transaction) { logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.') - return callback?.() - } - - for (const key in context) { - if (Object.hasOwn(context, key)) { - const value = context[key] - if (typeof value === 'object' || typeof value === 'function') { - logger.warn(`Invalid attribute type for ${key}. Skipped.`) - delete context[key] - } else if (key.indexOf('llm.') !== 0) { - logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`) - delete context[key] - context[`llm.${key}`] = value - } + return callback() + } + + for (const [key, value] of Object.entries(context)) { + if (typeof value === 'object' || typeof value === 'function') { + logger.warn(`Invalid attribute type for ${key}. Skipped.`) + delete context[key] + } else if (key.indexOf('llm.') !== 0) { + logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`) + delete context[key] + context[`llm.${key}`] = value } } diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js index e73b9e20f6..fb096f2c3d 100644 --- a/lib/util/llm-utils.js +++ b/lib/util/llm-utils.js @@ -29,13 +29,9 @@ function extractLlmAttribtues(context) { * @returns {Object} LLM context object */ function extractLlmContext(agent) { - let context = agent.tracer.getTransaction()._llmContextManager + const context = agent.tracer.getTransaction()._llmContextManager ? agent.tracer.getTransaction()._llmContextManager.getStore() : null - const llmContext = {} - while (context) { - Object.assign(llmContext, extractLlmAttribtues(context)) - context = null - } - return llmContext + + return extractLlmAttribtues(context) } diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index 772253fc2d..f3038bc836 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -128,6 +128,11 @@ tap.test('Agent API LLM methods', (t) => { return tx } + t.doesNotThrow(() => { + api.withLlmCustomAttributes(null, null) + t.equal(loggerMock.warn.callCount, 1) + }) + api.withLlmCustomAttributes( { 'toRename': 'value1', From 6e629f9616a9c297a7e7c98afcadf0e99c781a76 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 20 Aug 2024 10:44:35 -0700 Subject: [PATCH 26/33] fix: PR feedback --- api.js | 4 ++-- lib/util/llm-utils.js | 2 +- test/unit/util/llm-utils.test.js | 6 ++++++ 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/api.js b/api.js index d7bf0775b0..df08b65525 100644 --- a/api.js +++ b/api.js @@ -1904,7 +1904,7 @@ API.prototype.ignoreApdex = function ignoreApdex() { } /** - * Runs a function synchronously within a provided LLM custom attributes context and returns its return value. + * Run a function with the passed in LLM context as the active context and return its return value. * * An example of setting a custom attribute: * @@ -1912,7 +1912,7 @@ API.prototype.ignoreApdex = function ignoreApdex() { * return; * }) * @param {Object} context LLM custom attributes context - * @param {Function} callback synchronous function called within the context + * @param {Function} callback The function to execute in context. */ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context, callback) { const metric = this.agent.metrics.getOrCreateMetric( diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js index fb096f2c3d..5986f6455c 100644 --- a/lib/util/llm-utils.js +++ b/lib/util/llm-utils.js @@ -31,7 +31,7 @@ function extractLlmAttribtues(context) { function extractLlmContext(agent) { const context = agent.tracer.getTransaction()._llmContextManager ? agent.tracer.getTransaction()._llmContextManager.getStore() - : null + : {} return extractLlmAttribtues(context) } diff --git a/test/unit/util/llm-utils.test.js b/test/unit/util/llm-utils.test.js index 1e0ccc1c30..37dfe5194a 100644 --- a/test/unit/util/llm-utils.test.js +++ b/test/unit/util/llm-utils.test.js @@ -35,6 +35,12 @@ tap.test('extractLlmContext', (t) => { } } + tx._llmContextManager.run(null, () => { + const llmContext = extractLlmContext(agent) + t.equal(typeof llmContext, 'object') + t.equal(Object.entries(llmContext).length, 0) + }) + tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => { const llmContext = extractLlmContext(agent) t.equal(llmContext['llm.test'], 1) From 6d3bd2def85385c966d5f0ec2d3fc89008582822 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Tue, 20 Aug 2024 10:59:30 -0700 Subject: [PATCH 27/33] fix: PR feedback --- lib/util/llm-utils.js | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js index 5986f6455c..3b40793c69 100644 --- a/lib/util/llm-utils.js +++ b/lib/util/llm-utils.js @@ -29,9 +29,6 @@ function extractLlmAttribtues(context) { * @returns {Object} LLM context object */ function extractLlmContext(agent) { - const context = agent.tracer.getTransaction()._llmContextManager - ? agent.tracer.getTransaction()._llmContextManager.getStore() - : {} - + const context = agent.tracer.getTransaction()._llmContextManager?.getStore() || {} return extractLlmAttribtues(context) } From 1f2d240ad7ba320e21545b10704198b6ccbbc599 Mon Sep 17 00:00:00 2001 From: Webrealizer Date: Wed, 21 Aug 2024 10:02:11 -0700 Subject: [PATCH 28/33] Update lib/instrumentation/openai.js Co-authored-by: Bob Evans --- lib/instrumentation/openai.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index c85c9758e9..36e4487057 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -80,7 +80,7 @@ function recordEvent({ agent, type, msg }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign({}, msg, llmContext || {}) + Object.assign({}, msg, llmContext) ]) } From 8cc326f47e2cd8ffb52af97c67872819006e4960 Mon Sep 17 00:00:00 2001 From: Webrealizer Date: Wed, 21 Aug 2024 10:02:21 -0700 Subject: [PATCH 29/33] Update lib/instrumentation/langchain/common.js Co-authored-by: Bob Evans --- lib/instrumentation/langchain/common.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/instrumentation/langchain/common.js b/lib/instrumentation/langchain/common.js index 3644f77485..34e3d84c8e 100644 --- a/lib/instrumentation/langchain/common.js +++ b/lib/instrumentation/langchain/common.js @@ -54,7 +54,7 @@ common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign({}, msg, llmContext || {}) + Object.assign({}, msg, llmContext) ]) } From cab2f74b1239b888eecfeb4d2ef479d982776066 Mon Sep 17 00:00:00 2001 From: Webrealizer Date: Wed, 21 Aug 2024 10:02:30 -0700 Subject: [PATCH 30/33] Update lib/instrumentation/aws-sdk/v3/bedrock.js Co-authored-by: Bob Evans --- lib/instrumentation/aws-sdk/v3/bedrock.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 2d6a6d85ee..7fe1b4d277 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -60,7 +60,7 @@ function recordEvent({ agent, type, msg }) { agent.customEventAggregator.add([ { type, timestamp: Date.now() }, - Object.assign({}, msg, llmContext || {}) + Object.assign({}, msg, llmContext) ]) } From 67a35deb01590fa53204172c706516adfc4ed2ba Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 21 Aug 2024 10:05:16 -0700 Subject: [PATCH 31/33] fix: Improve test coverage --- test/unit/api/api-llm.test.js | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index f3038bc836..77c5aafeae 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -123,6 +123,15 @@ tap.test('Agent API LLM methods', (t) => { t.test('withLlmCustomAttributes', (t) => { const { api } = t.context + t.doesNotThrow(() => { + t.equal( + api.withLlmCustomAttributes({ test: 1 }, () => { + t.equal(loggerMock.warn.callCount, 1) + return 1 + }), + 1 + ) + }) helper.runInTransaction(api.agent, (tx) => { t.context.agent.tracer.getTransaction = () => { return tx @@ -130,7 +139,7 @@ tap.test('Agent API LLM methods', (t) => { t.doesNotThrow(() => { api.withLlmCustomAttributes(null, null) - t.equal(loggerMock.warn.callCount, 1) + t.equal(loggerMock.warn.callCount, 2) }) api.withLlmCustomAttributes( From cb423d6e838f457a71358fae6329f297ce91a985 Mon Sep 17 00:00:00 2001 From: mvazhenin Date: Wed, 21 Aug 2024 12:04:03 -0700 Subject: [PATCH 32/33] fix: More unit test and PR feedback --- api.js | 2 +- lib/util/llm-utils.js | 10 ++-- test/unit/api/api-llm.test.js | 88 +++++++++++++++++++++++++------- test/unit/util/llm-utils.test.js | 60 +++++++++++++++------- 4 files changed, 117 insertions(+), 43 deletions(-) diff --git a/api.js b/api.js index df08b65525..5cc7d59655 100644 --- a/api.js +++ b/api.js @@ -1932,7 +1932,7 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context return callback() } - for (const [key, value] of Object.entries(context)) { + for (const [key, value] of Object.entries(context || {})) { if (typeof value === 'object' || typeof value === 'function') { logger.warn(`Invalid attribute type for ${key}. Skipped.`) delete context[key] diff --git a/lib/util/llm-utils.js b/lib/util/llm-utils.js index 3b40793c69..7d26e692c7 100644 --- a/lib/util/llm-utils.js +++ b/lib/util/llm-utils.js @@ -5,7 +5,7 @@ 'use strict' -exports = module.exports = { extractLlmContext, extractLlmAttribtues } +exports = module.exports = { extractLlmContext, extractLlmAttributes } /** * Extract LLM attributes from the LLM context @@ -13,8 +13,8 @@ exports = module.exports = { extractLlmContext, extractLlmAttribtues } * @param {Object} context LLM context object * @returns {Object} LLM custom attributes */ -function extractLlmAttribtues(context) { - return Object.keys(context || {}).reduce((result, key) => { +function extractLlmAttributes(context) { + return Object.keys(context).reduce((result, key) => { if (key.indexOf('llm.') === 0) { result[key] = context[key] } @@ -29,6 +29,6 @@ function extractLlmAttribtues(context) { * @returns {Object} LLM context object */ function extractLlmContext(agent) { - const context = agent.tracer.getTransaction()._llmContextManager?.getStore() || {} - return extractLlmAttribtues(context) + const context = agent.tracer.getTransaction()?._llmContextManager?.getStore() || {} + return extractLlmAttributes(context) } diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index 77c5aafeae..b9baeea8f1 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -121,27 +121,58 @@ tap.test('Agent API LLM methods', (t) => { }) }) - t.test('withLlmCustomAttributes', (t) => { + t.test('withLlmCustomAttributes should handle no active transaction', (t) => { const { api } = t.context - t.doesNotThrow(() => { + t.autoend() + + t.equal( + api.withLlmCustomAttributes({ test: 1 }, () => { + t.equal(loggerMock.warn.callCount, 1) + return 1 + }), + 1 + ) + }) + + t.test('withLlmCustomAttributes should handle an empty store', (t) => { + const { api } = t.context + const agent = api.agent + + t.autoend() + helper.runInTransaction(api.agent, (tx) => { + agent.tracer.getTransaction = () => { + return tx + } t.equal( - api.withLlmCustomAttributes({ test: 1 }, () => { - t.equal(loggerMock.warn.callCount, 1) + api.withLlmCustomAttributes(null, () => { return 1 }), 1 ) }) + }) + + t.test('withLlmCustomAttributes should handle no callback', (t) => { + const { api } = t.context + const agent = api.agent + t.autoend() helper.runInTransaction(api.agent, (tx) => { - t.context.agent.tracer.getTransaction = () => { + agent.tracer.getTransaction = () => { return tx } + api.withLlmCustomAttributes({ test: 1 }, null) + t.equal(loggerMock.warn.callCount, 1) + }) + }) - t.doesNotThrow(() => { - api.withLlmCustomAttributes(null, null) - t.equal(loggerMock.warn.callCount, 2) - }) - + t.test('withLlmCustomAttributes should normalize attributes', (t) => { + const { api } = t.context + const agent = api.agent + t.autoend() + helper.runInTransaction(api.agent, (tx) => { + agent.tracer.getTransaction = () => { + return tx + } api.withLlmCustomAttributes( { 'toRename': 'value1', @@ -160,19 +191,40 @@ tap.test('Agent API LLM methods', (t) => { t.notOk(parentContext.toDelete3) t.equal(parentContext['llm.number'], 1) t.equal(parentContext['llm.boolean'], true) - - api.withLlmCustomAttributes({ 'llm.someAttribute': 'someValue' }, () => { - const contextManager = tx._llmContextManager - const context = contextManager.getStore() - t.equal(context[`llm.toRename`], 'value1') - t.equal(context['llm.someAttribute'], 'someValue') - t.end() - }) } ) }) }) + t.test('withLlmCustomAttributes should support branching', (t) => { + const { api } = t.context + const agent = api.agent + t.autoend() + helper.runInTransaction(api.agent, (tx) => { + agent.tracer.getTransaction = () => { + return tx + } + api.withLlmCustomAttributes({ 'llm.step': '1', 'llm.path': 'root' }, () => { + const contextManager = tx._llmContextManager + const context = contextManager.getStore() + t.equal(context[`llm.step`], '1') + t.equal(context['llm.path'], 'root') + api.withLlmCustomAttributes({ 'llm.step': '1.1', 'llm.path': 'root/1' }, () => { + const contextManager = tx._llmContextManager + const context = contextManager.getStore() + t.equal(context[`llm.step`], '1.1') + t.equal(context['llm.path'], 'root/1') + }) + api.withLlmCustomAttributes({ 'llm.step': '1.2', 'llm.path': 'root/2' }, () => { + const contextManager = tx._llmContextManager + const context = contextManager.getStore() + t.equal(context[`llm.step`], '1.2') + t.equal(context['llm.path'], 'root/2') + }) + }) + }) + }) + t.test('setLlmTokenCount should register callback to calculate token counts', async (t) => { const { api, agent } = t.context function callback(model, content) { diff --git a/test/unit/util/llm-utils.test.js b/test/unit/util/llm-utils.test.js index 37dfe5194a..6ce2bf40d6 100644 --- a/test/unit/util/llm-utils.test.js +++ b/test/unit/util/llm-utils.test.js @@ -6,7 +6,7 @@ 'use strict' const tap = require('tap') -const { extractLlmAttribtues, extractLlmContext } = require('../../../lib/util/llm-utils') +const { extractLlmAttributes, extractLlmContext } = require('../../../lib/util/llm-utils') const { AsyncLocalStorage } = require('async_hooks') tap.test('extractLlmAttributes', (t) => { @@ -16,7 +16,7 @@ tap.test('extractLlmAttributes', (t) => { 'fllm.skip': 3 } - const llmContext = extractLlmAttribtues(context) + const llmContext = extractLlmAttributes(context) t.notOk(llmContext.skip) t.notOk(llmContext['fllm.skip']) t.equal(llmContext['llm.get'], 2) @@ -24,27 +24,49 @@ tap.test('extractLlmAttributes', (t) => { }) tap.test('extractLlmContext', (t) => { - const tx = { - _llmContextManager: new AsyncLocalStorage() - } - const agent = { - tracer: { - getTransaction: () => { - return tx + let tx + let agent + t.autoend() + t.beforeEach(() => { + tx = { + _llmContextManager: new AsyncLocalStorage() + } + agent = { + tracer: { + getTransaction: () => { + return tx + } } } - } + }) - tx._llmContextManager.run(null, () => { - const llmContext = extractLlmContext(agent) - t.equal(typeof llmContext, 'object') - t.equal(Object.entries(llmContext).length, 0) + t.test('handle empty context', (t) => { + t.autoend() + tx._llmContextManager.run(null, () => { + const llmContext = extractLlmContext(agent) + t.equal(typeof llmContext, 'object') + t.equal(Object.entries(llmContext).length, 0) + }) }) - tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => { - const llmContext = extractLlmContext(agent) - t.equal(llmContext['llm.test'], 1) - t.notOk(llmContext.skip) - t.end() + t.test('extract LLM context', (t) => { + t.autoend() + tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => { + const llmContext = extractLlmContext(agent) + t.equal(llmContext['llm.test'], 1) + t.notOk(llmContext.skip) + }) + }) + + t.test('no transaction', (t) => { + t.autoend() + agent.tracer.getTransaction = () => { + return null + } + tx._llmContextManager.run(null, () => { + const llmContext = extractLlmContext(agent) + t.equal(typeof llmContext, 'object') + t.equal(Object.entries(llmContext).length, 0) + }) }) }) From 21e4b5f39c9fbefcae4c014da7cafbc01a8fe7ac Mon Sep 17 00:00:00 2001 From: Bob Evans Date: Thu, 22 Aug 2024 16:26:10 -0400 Subject: [PATCH 33/33] chore: Addressed code review feedback --- api.js | 7 +-- test/unit/api/api-llm.test.js | 48 +++++++++++-------- test/unit/util/llm-utils.test.js | 20 ++++---- test/versioned/openai/chat-completions.tap.js | 34 ------------- 4 files changed, 42 insertions(+), 67 deletions(-) diff --git a/api.js b/api.js index 5cc7d59655..49c210f54e 100644 --- a/api.js +++ b/api.js @@ -1915,6 +1915,7 @@ API.prototype.ignoreApdex = function ignoreApdex() { * @param {Function} callback The function to execute in context. */ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context, callback) { + context = context || {} const metric = this.agent.metrics.getOrCreateMetric( NAMES.SUPPORTABILITY.API + '/withLlmCustomAttributes' ) @@ -1932,7 +1933,7 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context return callback() } - for (const [key, value] of Object.entries(context || {})) { + for (const [key, value] of Object.entries(context)) { if (typeof value === 'object' || typeof value === 'function') { logger.warn(`Invalid attribute type for ${key}. Skipped.`) delete context[key] @@ -1944,9 +1945,9 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context } transaction._llmContextManager = transaction._llmContextManager || new AsyncLocalStorage() - const parentContext = transaction._llmContextManager.getStore() + const parentContext = transaction._llmContextManager.getStore() || {} - const fullContext = Object.assign({}, parentContext || {}, context || {}) + const fullContext = Object.assign({}, parentContext, context) return transaction._llmContextManager.run(fullContext, callback) } diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js index b9baeea8f1..26eab92c81 100644 --- a/test/unit/api/api-llm.test.js +++ b/test/unit/api/api-llm.test.js @@ -123,8 +123,6 @@ tap.test('Agent API LLM methods', (t) => { t.test('withLlmCustomAttributes should handle no active transaction', (t) => { const { api } = t.context - t.autoend() - t.equal( api.withLlmCustomAttributes({ test: 1 }, () => { t.equal(loggerMock.warn.callCount, 1) @@ -132,13 +130,13 @@ tap.test('Agent API LLM methods', (t) => { }), 1 ) + t.end() }) t.test('withLlmCustomAttributes should handle an empty store', (t) => { const { api } = t.context const agent = api.agent - t.autoend() helper.runInTransaction(api.agent, (tx) => { agent.tracer.getTransaction = () => { return tx @@ -149,26 +147,26 @@ tap.test('Agent API LLM methods', (t) => { }), 1 ) + t.end() }) }) t.test('withLlmCustomAttributes should handle no callback', (t) => { const { api } = t.context const agent = api.agent - t.autoend() helper.runInTransaction(api.agent, (tx) => { agent.tracer.getTransaction = () => { return tx } api.withLlmCustomAttributes({ test: 1 }, null) t.equal(loggerMock.warn.callCount, 1) + t.end() }) }) t.test('withLlmCustomAttributes should normalize attributes', (t) => { const { api } = t.context const agent = api.agent - t.autoend() helper.runInTransaction(api.agent, (tx) => { agent.tracer.getTransaction = () => { return tx @@ -191,6 +189,7 @@ tap.test('Agent API LLM methods', (t) => { t.notOk(parentContext.toDelete3) t.equal(parentContext['llm.number'], 1) t.equal(parentContext['llm.boolean'], true) + t.end() } ) }) @@ -204,24 +203,31 @@ tap.test('Agent API LLM methods', (t) => { agent.tracer.getTransaction = () => { return tx } - api.withLlmCustomAttributes({ 'llm.step': '1', 'llm.path': 'root' }, () => { - const contextManager = tx._llmContextManager - const context = contextManager.getStore() - t.equal(context[`llm.step`], '1') - t.equal(context['llm.path'], 'root') - api.withLlmCustomAttributes({ 'llm.step': '1.1', 'llm.path': 'root/1' }, () => { - const contextManager = tx._llmContextManager - const context = contextManager.getStore() - t.equal(context[`llm.step`], '1.1') - t.equal(context['llm.path'], 'root/1') - }) - api.withLlmCustomAttributes({ 'llm.step': '1.2', 'llm.path': 'root/2' }, () => { + api.withLlmCustomAttributes( + { 'llm.step': '1', 'llm.path': 'root', 'llm.name': 'root' }, + () => { const contextManager = tx._llmContextManager const context = contextManager.getStore() - t.equal(context[`llm.step`], '1.2') - t.equal(context['llm.path'], 'root/2') - }) - }) + t.equal(context[`llm.step`], '1') + t.equal(context['llm.path'], 'root') + t.equal(context['llm.name'], 'root') + api.withLlmCustomAttributes({ 'llm.step': '1.1', 'llm.path': 'root/1' }, () => { + const contextManager = tx._llmContextManager + const context = contextManager.getStore() + t.equal(context[`llm.step`], '1.1') + t.equal(context['llm.path'], 'root/1') + t.equal(context['llm.name'], 'root') + }) + api.withLlmCustomAttributes({ 'llm.step': '1.2', 'llm.path': 'root/2' }, () => { + const contextManager = tx._llmContextManager + const context = contextManager.getStore() + t.equal(context[`llm.step`], '1.2') + t.equal(context['llm.path'], 'root/2') + t.equal(context['llm.name'], 'root') + t.end() + }) + } + ) }) }) diff --git a/test/unit/util/llm-utils.test.js b/test/unit/util/llm-utils.test.js index 6ce2bf40d6..666a42e3f2 100644 --- a/test/unit/util/llm-utils.test.js +++ b/test/unit/util/llm-utils.test.js @@ -24,42 +24,42 @@ tap.test('extractLlmAttributes', (t) => { }) tap.test('extractLlmContext', (t) => { - let tx - let agent - t.autoend() - t.beforeEach(() => { - tx = { + t.beforeEach((t) => { + const tx = { _llmContextManager: new AsyncLocalStorage() } - agent = { + t.context.agent = { tracer: { getTransaction: () => { return tx } } } + t.context.tx = tx }) t.test('handle empty context', (t) => { - t.autoend() + const { tx, agent } = t.context tx._llmContextManager.run(null, () => { const llmContext = extractLlmContext(agent) t.equal(typeof llmContext, 'object') t.equal(Object.entries(llmContext).length, 0) + t.end() }) }) t.test('extract LLM context', (t) => { - t.autoend() + const { tx, agent } = t.context tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => { const llmContext = extractLlmContext(agent) t.equal(llmContext['llm.test'], 1) t.notOk(llmContext.skip) + t.end() }) }) t.test('no transaction', (t) => { - t.autoend() + const { tx, agent } = t.context agent.tracer.getTransaction = () => { return null } @@ -67,6 +67,8 @@ tap.test('extractLlmContext', (t) => { const llmContext = extractLlmContext(agent) t.equal(typeof llmContext, 'object') t.equal(Object.entries(llmContext).length, 0) + t.end() }) }) + t.end() }) diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js index 054a4a7231..322691bebb 100644 --- a/test/versioned/openai/chat-completions.tap.js +++ b/test/versioned/openai/chat-completions.tap.js @@ -449,40 +449,6 @@ tap.test('OpenAI instrumentation - chat completions', (t) => { }) }) - t.test('should create chat completion message and summary for every message sent', (test) => { - const { client, agent } = t.context - helper.runInTransaction(agent, async (tx) => { - const model = 'gpt-3.5-turbo-0613' - const content = 'You are a mathematician.' - await client.chat.completions.create({ - max_tokens: 100, - temperature: 0.5, - model, - messages: [ - { role: 'user', content }, - { role: 'user', content: 'What does 1 plus 1 equal?' } - ] - }) - - const events = agent.customEventAggregator.events.toArray() - test.equal(events.length, 4, 'should create a chat completion message and summary event') - const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage') - test.llmMessages({ - tx, - chatMsgs, - model, - id: 'chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat', - resContent: '1 plus 2 is 3.', - reqContent: content - }) - - const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0] - test.llmSummary({ tx, model, chatSummary, tokenUsage: true }) - tx.end() - test.end() - }) - }) - t.test('should record LLM custom events with attributes', (test) => { const { client, agent } = t.context const api = helper.getAgentApi()