feat: Added ability to disable capturing of content in LLM events `ai…

…_monitoring.record_content.enabled` (newrelic#268)
jsumners-nr · Mar 4, 2024 · a36050e · a36050e
1 parent e5565f2
commit a36050e
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 6 deletions.
diff --git a/lib/llm/chat-completion-message.js b/lib/llm/chat-completion-message.js
@@ -39,20 +39,22 @@ class LlmChatCompletionMessage extends LlmEvent {
     params = Object.assign({}, defaultParams, params)
     super(params)
 
-    const { content, isResponse, index, completionId } = params
+    const { agent, content, isResponse, index, completionId } = params
 
     this.is_response = isResponse
     this.completion_id = completionId
     this.sequence = index
-    this.content = content
+    this.content = agent.config?.ai_monitoring?.record_content?.enabled ? content : undefined
     this.role = ''
 
     this.#setId(index)
     if (this.is_response === true) {
       this.role = 'assistant'
     } else {
       this.role = 'user'
-      this.content = this.bedrockCommand.prompt
+      this.content = agent.config?.ai_monitoring?.record_content?.enabled
+        ? this.bedrockCommand.prompt
+        : undefined
     }
   }
 

diff --git a/lib/llm/embedding.js b/lib/llm/embedding.js
@@ -19,8 +19,11 @@ const defaultParams = {}
 class LlmEmbedding extends LlmEvent {
   constructor(params = defaultParams) {
     super(params)
+    const { agent } = params
 
-    this.input = this.bedrockCommand.prompt
+    this.input = agent.config?.ai_monitoring?.record_content?.enabled
+      ? this.bedrockCommand.prompt
+      : undefined
     this.error = params.isError
     this.duration = params.segment.getDurationInMillis()
     this['response.usage.total_tokens'] = this.bedrockResponse.inputTokenCount

diff --git a/tests/unit/llm/chat-completion-message.tap.js b/tests/unit/llm/chat-completion-message.tap.js
@@ -16,6 +16,12 @@ tap.beforeEach((t) => {
     config: {
       applications() {
         return ['test-app']
+      },
+      ai_monitoring: {
+        enabled: true,
+        record_content: {
+          enabled: true
+        }
       }
     },
     tracer: {
@@ -126,3 +132,13 @@ tap.test('create creates a ai21 response instance when response.id is undefined'
   t.equal(event.role, 'assistant')
   t.match(event.id, /[\w-]{36}-0/)
 })
+
+tap.test(
+  'should not capture content when `ai_monitoring.record_content.enabled` is false',
+  async (t) => {
+    const { agent } = t.context
+    agent.config.ai_monitoring.record_content.enabled = false
+    const event = new LlmChatCompletionMessage(t.context)
+    t.equal(event.content, undefined, 'content should be empty')
+  }
+)
diff --git a/tests/unit/llm/embedding.tap.js b/tests/unit/llm/embedding.tap.js
@@ -16,6 +16,12 @@ tap.beforeEach((t) => {
     config: {
       applications() {
         return ['test-app']
+      },
+      ai_monitoring: {
+        enabled: true,
+        record_content: {
+          enabled: true
+        }
       }
     },
     tracer: {
@@ -61,3 +67,13 @@ tap.test('creates a basic embedding', async (t) => {
   t.equal(event['response.usage.total_tokens'], 0)
   t.equal(event['response.usage.prompt_tokens'], 0)
 })
+
+tap.test(
+  'should not capture input when `ai_monitoring.record_content.enabled` is false',
+  async (t) => {
+    const { agent } = t.context
+    agent.config.ai_monitoring.record_content.enabled = false
+    const event = new LlmEmbedding(t.context)
+    t.equal(event.input, undefined, 'input should be empty')
+  }
+)
diff --git a/tests/versioned/v3/bedrock-chat-completions.tap.js b/tests/versioned/v3/bedrock-chat-completions.tap.js
@@ -43,13 +43,16 @@ const requests = {
 tap.beforeEach(async (t) => {
   const helper = utils.TestAgent.makeInstrumented({
     ai_monitoring: {
-      enabled: true
+      enabled: true,
+      streaming: {
+        enabled: true
+      }
     }
   })
   common.registerInstrumentation(helper)
   // TODO: cannot set in config as configuration does not exist on agent
   // just yet. update this when agent is released
-  helper.agent.config.ai_monitoring.streaming = { enabled: true }
+  helper.agent.config.ai_monitoring.record_content = { enabled: true }
   t.context.helper = helper
 
   const bedrock = require('@aws-sdk/client-bedrock-runtime')

diff --git a/tests/versioned/v3/bedrock-embeddings.tap.js b/tests/versioned/v3/bedrock-embeddings.tap.js
@@ -30,6 +30,9 @@ tap.beforeEach(async (t) => {
     }
   })
   common.registerInstrumentation(helper)
+  // TODO: cannot set in config as configuration does not exist on agent
+  // just yet. update this when agent is released
+  helper.agent.config.ai_monitoring.record_content = { enabled: true }
   t.context.helper = helper
 
   const bedrock = require('@aws-sdk/client-bedrock-runtime')