From 7069335bfee38b1774da00bdbb63138ebf38da90 Mon Sep 17 00:00:00 2001 From: James Sumners Date: Tue, 9 Jul 2024 15:30:33 -0400 Subject: [PATCH] feat: Added support for account level governance of AI Monitoring (#2326) --- lib/config/index.js | 14 ++++++++--- lib/instrumentation/aws-sdk/v3/bedrock.js | 25 +++++++++++++++++-- lib/instrumentation/langchain/runnable.js | 14 ++++++++++- lib/instrumentation/langchain/tools.js | 11 +++++++- lib/instrumentation/langchain/vectorstore.js | 18 ++++++++++--- lib/instrumentation/openai.js | 23 +++++++++++++++-- test/unit/config/config-server-side.test.js | 9 +++++++ .../langchain/runnables.test.js | 2 +- 8 files changed, 101 insertions(+), 15 deletions(-) diff --git a/lib/config/index.js b/lib/config/index.js index c5ba7ddfc9..62c77d1706 100644 --- a/lib/config/index.js +++ b/lib/config/index.js @@ -306,6 +306,12 @@ Config.prototype._fromServer = function _fromServer(params, key) { case 'high_security': break + // interpret AI Monitoring account setting + case 'collect_ai': + this._disableOption(params.collect_ai, 'ai_monitoring') + this.emit('change', this) + break + // always accept these settings case 'cross_process_id': case 'encoding_key': @@ -1071,10 +1077,10 @@ function setFromEnv({ config, key, envVar, formatter, paths }) { /** * Recursively visit the nodes of the config definition and look for environment variable names, overriding any configuration values that are found. * - * @param {object} [config=this] The current level of the configuration object. - * @param {object} [data=configDefinition] The current level of the config definition object. - * @param {Array} [paths=[]] keeps track of the nested path to properly derive the env var - * @param {number} [objectKeys=1] indicator of how many keys exist in current node to know when to remove current node after all keys are processed + * @param {object} [config] The current level of the configuration object. + * @param {object} [data] The current level of the config definition object. + * @param {Array} [paths] keeps track of the nested path to properly derive the env var + * @param {number} [objectKeys] indicator of how many keys exist in current node to know when to remove current node after all keys are processed */ Config.prototype._fromEnvironment = function _fromEnvironment( config = this, diff --git a/lib/instrumentation/aws-sdk/v3/bedrock.js b/lib/instrumentation/aws-sdk/v3/bedrock.js index 9fb575aad2..055c572c63 100644 --- a/lib/instrumentation/aws-sdk/v3/bedrock.js +++ b/lib/instrumentation/aws-sdk/v3/bedrock.js @@ -85,8 +85,21 @@ function addLlmMeta({ agent, segment }) { * @param {BedrockCommand} params.bedrockCommand parsed input * @param {Error|null} params.err error from request if exists * @param params.bedrockResponse + * @param params.shim */ -function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockResponse, err }) { +function recordChatCompletionMessages({ + agent, + shim, + segment, + bedrockCommand, + bedrockResponse, + err +}) { + if (shouldSkipInstrumentation(agent.config) === true) { + shim.logger.debug('skipping sending of ai data') + return + } + const summary = new LlmChatCompletionSummary({ agent, bedrockResponse, @@ -133,12 +146,18 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR * * @param {object} params function params * @param {object} params.agent instance of agent + * @param {object} params.shim current shim instance * @param {object} params.segment active segment * @param {BedrockCommand} params.bedrockCommand parsed input * @param {Error|null} params.err error from request if exists * @param params.bedrockResponse */ -function recordEmbeddingMessage({ agent, segment, bedrockCommand, bedrockResponse, err }) { +function recordEmbeddingMessage({ agent, shim, segment, bedrockCommand, bedrockResponse, err }) { + if (shouldSkipInstrumentation(agent.config) === true) { + shim.logger.debug('skipping sending of ai data') + return + } + const embedding = new LlmEmbedding({ agent, segment, @@ -239,6 +258,7 @@ function handleResponse({ shim, err, response, segment, bedrockCommand, modelTyp if (modelType === 'completion') { recordChatCompletionMessages({ agent, + shim, segment, bedrockCommand, bedrockResponse, @@ -247,6 +267,7 @@ function handleResponse({ shim, err, response, segment, bedrockCommand, modelTyp } else if (modelType === 'embedding') { recordEmbeddingMessage({ agent, + shim, segment, bedrockCommand, bedrockResponse, diff --git a/lib/instrumentation/langchain/runnable.js b/lib/instrumentation/langchain/runnable.js index 2bfbd90ee6..b6d76f4ba6 100644 --- a/lib/instrumentation/langchain/runnable.js +++ b/lib/instrumentation/langchain/runnable.js @@ -16,13 +16,14 @@ const LlmErrorMessage = require('../../llm-events/error-message') const { DESTINATIONS } = require('../../config/attribute-filter') const { langchainRunId } = require('../../symbols') const { RecorderSpec } = require('../../shim/specs') +const { shouldSkipInstrumentation } = require('./common') module.exports = function initialize(shim, langchain) { const { agent, pkgVersion } = shim if (common.shouldSkipInstrumentation(agent.config)) { shim.logger.debug( - 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' + 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' ) return } @@ -186,6 +187,15 @@ function wrapNextHandler({ shim, output, segment, request, metadata, tags }) { function recordChatCompletionEvents({ segment, messages, events, metadata, tags, err, shim }) { const { pkgVersion, agent } = shim segment.end() + + if (shouldSkipInstrumentation(shim.agent.config) === true) { + // We need this check inside the wrapper because it is possible for monitoring + // to be disabled at the account level. In such a case, the value is set + // after the instrumentation has been initialized. + shim.logger.debug('skipping sending of ai data') + return + } + const completionSummary = new LangChainCompletionSummary({ agent, messages, @@ -198,6 +208,7 @@ function recordChatCompletionEvents({ segment, messages, events, metadata, tags, common.recordEvent({ agent, + shim, type: 'LlmChatCompletionSummary', pkgVersion, msg: completionSummary @@ -266,6 +277,7 @@ function recordCompletions({ events, completionSummary, agent, segment, shim }) common.recordEvent({ agent, + shim, type: 'LlmChatCompletionMessage', pkgVersion: shim.pkgVersion, msg: completionMsg diff --git a/lib/instrumentation/langchain/tools.js b/lib/instrumentation/langchain/tools.js index 9844f6b3cf..17b0178998 100644 --- a/lib/instrumentation/langchain/tools.js +++ b/lib/instrumentation/langchain/tools.js @@ -35,6 +35,15 @@ module.exports = function initialize(shim, tools) { const metadata = mergeMetadata(instanceMeta, paramsMeta) const tags = mergeTags(instanceTags, paramsTags) segment.end() + + if (shouldSkipInstrumentation(shim.agent.config) === true) { + // We need this check inside the wrapper because it is possible for monitoring + // to be disabled at the account level. In such a case, the value is set + // after the instrumentation has been initialized. + shim.logger.debug('skipping sending of ai data') + return + } + const toolEvent = new LangChainTool({ agent, description, @@ -47,7 +56,7 @@ module.exports = function initialize(shim, tools) { segment, error: err != null }) - recordEvent({ agent, type: 'LlmTool', pkgVersion, msg: toolEvent }) + recordEvent({ agent, shim, type: 'LlmTool', pkgVersion, msg: toolEvent }) if (err) { agent.errors.add( diff --git a/lib/instrumentation/langchain/vectorstore.js b/lib/instrumentation/langchain/vectorstore.js index d602b9a169..61d27174b0 100644 --- a/lib/instrumentation/langchain/vectorstore.js +++ b/lib/instrumentation/langchain/vectorstore.js @@ -23,11 +23,12 @@ const LlmErrorMessage = require('../../llm-events/error-message') * @param {number} params.k vector search top k * @param {object} params.output vector search documents * @param {Agent} params.agent NR agent instance + * @param {Shim} params.shim current shim instance * @param {TraceSegment} params.segment active segment from vector search * @param {string} params.pkgVersion langchain version - * @param {err} params.err if it exists + * @param {Error} params.err if it exists */ -function recordVectorSearch({ request, k, output, agent, segment, pkgVersion, err }) { +function recordVectorSearch({ request, k, output, agent, shim, segment, pkgVersion, err }) { const vectorSearch = new LangChainVectorSearch({ agent, segment, @@ -37,7 +38,7 @@ function recordVectorSearch({ request, k, output, agent, segment, pkgVersion, er error: err !== null }) - recordEvent({ agent, type: 'LlmVectorSearch', pkgVersion, msg: vectorSearch }) + recordEvent({ agent, shim, type: 'LlmVectorSearch', pkgVersion, msg: vectorSearch }) output.forEach((document, sequence) => { const vectorSearchResult = new LangChainVectorSearchResult({ @@ -51,6 +52,7 @@ function recordVectorSearch({ request, k, output, agent, segment, pkgVersion, er recordEvent({ agent, + shim, type: 'LlmVectorSearchResult', pkgVersion, msg: vectorSearchResult @@ -97,7 +99,15 @@ module.exports = function initialize(shim, vectorstores) { } segment.end() - recordVectorSearch({ request, k, output, agent, segment, pkgVersion, err }) + if (shouldSkipInstrumentation(shim.agent.config) === true) { + // We need this check inside the wrapper because it is possible for monitoring + // to be disabled at the account level. In such a case, the value is set + // after the instrumentation has been initialized. + shim.logger.debug('skipping sending of ai data') + return + } + + recordVectorSearch({ request, k, output, agent, shim, segment, pkgVersion, err }) segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true) } diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js index a67196ccaf..da0b4a22f4 100644 --- a/lib/instrumentation/openai.js +++ b/lib/instrumentation/openai.js @@ -99,12 +99,18 @@ function addLlmMeta({ agent, segment }) { * * @param {object} params input params * @param {Agent} params.agent NR agent instance + * @param {Shim} params.shim the current shim instance * @param {TraceSegment} params.segment active segment from chat completion * @param {object} params.request chat completion params * @param {object} params.response chat completion response * @param {boolean} [params.err] err if it exists */ -function recordChatCompletionMessages({ agent, segment, request, response, err }) { +function recordChatCompletionMessages({ agent, shim, segment, request, response, err }) { + if (shouldSkipInstrumentation(agent.config, shim) === true) { + shim.logger.debug('skipping sending of ai data') + return + } + if (!response) { // If we get an error, it is possible that `response = null`. // In that case, we define it to be an empty object. @@ -195,6 +201,7 @@ function instrumentStream({ agent, shim, request, response, segment }) { recordChatCompletionMessages({ agent: shim.agent, + shim, segment, request, response: chunk, @@ -205,6 +212,7 @@ function instrumentStream({ agent, shim, request, response, segment }) { }) } +/* eslint-disable sonarjs/cognitive-complexity */ module.exports = function initialize(agent, openai, moduleName, shim) { if (shouldSkipInstrumentation(agent.config, shim)) { shim.logger.debug( @@ -268,6 +276,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) { } else { recordChatCompletionMessages({ agent, + shim, segment, request, response, @@ -301,10 +310,20 @@ module.exports = function initialize(agent, openai, moduleName, shim) { // In that case, we define it to be an empty object. response = {} } + + segment.end() + if (shouldSkipInstrumentation(shim.agent.config, shim) === true) { + // We need this check inside the wrapper because it is possible for monitoring + // to be disabled at the account level. In such a case, the value is set + // after the instrumentation has been initialized. + shim.logger.debug('skipping sending of ai data') + return + } + response.headers = segment[openAiHeaders] // explicitly end segment to get consistent duration // for both LLM events and the segment - segment.end() + const embedding = new LlmEmbedding({ agent, segment, diff --git a/test/unit/config/config-server-side.test.js b/test/unit/config/config-server-side.test.js index 4e47b17a41..b527e7b51d 100644 --- a/test/unit/config/config-server-side.test.js +++ b/test/unit/config/config-server-side.test.js @@ -161,6 +161,15 @@ tap.test('when receiving server-side configuration', (t) => { t.end() }) + t.test('should disable ai monitoring', (t) => { + config.ai_monitoring.enabled = true + t.equal(config.ai_monitoring.enabled, true) + config.onConnect({ collect_ai: false }) + t.equal(config.ai_monitoring.enabled, false) + + t.end() + }) + t.test('should configure cross application tracing', (t) => { config.cross_application_tracer.enabled = true diff --git a/test/unit/instrumentation/langchain/runnables.test.js b/test/unit/instrumentation/langchain/runnables.test.js index aaba0faa79..319c456b57 100644 --- a/test/unit/instrumentation/langchain/runnables.test.js +++ b/test/unit/instrumentation/langchain/runnables.test.js @@ -46,7 +46,7 @@ test('langchain/core/runnables unit tests', (t) => { t.equal(shim.logger.debug.callCount, 1, 'should log 1 debug messages') t.equal( shim.logger.debug.args[0][0], - 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' + 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' ) const isWrapped = shim.isWrapped(MockRunnable.RunnableSequence.prototype.invoke) t.equal(isWrapped, false, 'should not wrap runnable invoke')