feat: Added support for account level governance of AI Monitoring

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,

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,

lib/instrumentation/langchain/runnable.js

@@ -17,12 +17,14 @@ const { DESTINATIONS } = require('../../config/attribute-filter')
 const { langchainRunId } = require('../../symbols')
 const { RecorderSpec } = require('../../shim/specs')
 
+const skipRecordMessage = 'skipping sending of ai data'
+
 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
   }
@@ -54,6 +56,11 @@ function instrumentInvokeChain({ langchain, shim }) {
     langchain.RunnableSequence.prototype,
     'invoke',
     function wrapCall(shim, _invoke, fnName, args) {
+      if (common.shouldSkipInstrumentation(shim.agent.config) === true) {
+        shim.logger.debug(skipRecordMessage)
+        return
+      }
+
       const [request, params] = args
       const metadata = params?.metadata ?? {}
       const tags = params?.tags ?? []
@@ -89,6 +96,11 @@ function instrumentStream({ langchain, shim }) {
     langchain.RunnableSequence.prototype,
     'stream',
     function wrapStream(shim, _stream, fnName, args) {
+      if (common.shouldSkipInstrumentation(shim.agent.config) === true) {
+        shim.logger.debug(skipRecordMessage)
+        return
+      }
+
       const [request, params] = args
       const metadata = params?.metadata ?? {}
       const tags = params?.tags ?? []
@@ -132,6 +144,11 @@ function instrumentStream({ langchain, shim }) {
  * @param {Array} params.tags tags for the call
  */
 function wrapNextHandler({ shim, output, segment, request, metadata, tags }) {
+  if (common.shouldSkipInstrumentation(shim.agent.config) === true) {
+    shim.logger.debug(skipRecordMessage)
+    return
+  }
+
   shim.wrap(output, 'next', function wrapIterator(shim, orig) {
     let content = ''
     return async function wrappedIterator() {
@@ -198,6 +215,7 @@ function recordChatCompletionEvents({ segment, messages, events, metadata, tags,
 
   common.recordEvent({
     agent,
+    shim,
     type: 'LlmChatCompletionSummary',
     pkgVersion,
     msg: completionSummary
@@ -266,6 +284,7 @@ function recordCompletions({ events, completionSummary, agent, segment, shim })
 
     common.recordEvent({
       agent,
+      shim,
       type: 'LlmChatCompletionMessage',
       pkgVersion: shim.pkgVersion,
       msg: completionMsg

lib/instrumentation/langchain/tools.js

@@ -25,6 +25,14 @@ module.exports = function initialize(shim, tools) {
   }
 
   shim.record(tools.StructuredTool.prototype, 'call', function wrapCall(shim, call, fnName, args) {
+    if (shouldSkipInstrumentation(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 { name, metadata: instanceMeta, description, tags: instanceTags } = this
     const [request, params] = args
     const { metadata: paramsMeta, tags: paramsTags } = params || {}
@@ -47,7 +55,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(

lib/instrumentation/langchain/vectorstore.js

@@ -23,11 +23,13 @@ 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 params.shim
  */
-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 +39,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 +53,7 @@ function recordVectorSearch({ request, k, output, agent, segment, pkgVersion, er
 
     recordEvent({
       agent,
+      shim,
       type: 'LlmVectorSearchResult',
       pkgVersion,
       msg: vectorSearchResult
@@ -84,6 +87,14 @@ module.exports = function initialize(shim, vectorstores) {
     vectorstores.VectorStore.prototype,
     'similaritySearch',
     function wrapCall(shim, similaritySearch, fnName, args) {
+      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 [request, k] = args
 
       return new RecorderSpec({
@@ -97,7 +108,7 @@ module.exports = function initialize(shim, vectorstores) {
           }
 
           segment.end()
-          recordVectorSearch({ request, k, output, agent, segment, pkgVersion, err })
+          recordVectorSearch({ request, k, output, agent, shim, segment, pkgVersion, err })
 
           segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
         }

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,
@@ -289,6 +298,11 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
     openai.Embeddings.prototype,
     'create',
     function wrapEmbeddingCreate(shim, embeddingCreate, name, args) {
+      if (shouldSkipInstrumentation(agent.config, shim) === true) {
+        shim.logger.debug('skipping sending of ai data')
+        return
+      }
+
       const [request] = args
       return new RecorderSpec({
         name: OPENAI.EMBEDDING,

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
 

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')