Use AiMonitoringSubscriber abstraction

Author: amychisholm03

lib/llm-events/error-message.js

@@ -21,10 +21,10 @@ module.exports = class LlmErrorMessage {
    * @param {LlmVectorStoreSearch} [params.vectorsearch] Details about the vector
    * search if it was a vector search event.
    * @param {LlmTool} [params.tool] Details about the tool event if it was a tool event.
-   * @param {object} [params.agent] Details about the AI agent event if it was an AI agent event.
+   * @param {object} [params.aiAgent] Details about the AI agent event if it was an AI agent event.
    * @param {boolean} [params.useNameAsCode] defaults to false, only Bedrock sets it to true so far
    */
-  constructor({ response, cause, summary = {}, embedding = {}, vectorsearch = {}, tool = {}, agent = {}, useNameAsCode = false } = {}) {
+  constructor({ response, cause, summary = {}, embedding = {}, vectorsearch = {}, tool = {}, aiAgent = {}, useNameAsCode = false } = {}) {
     this['http.statusCode'] = response?.statusCode ?? response?.status ?? cause?.status
     this['error.message'] = cause?.message
     this['error.code'] = response?.code ?? cause?.error?.code
@@ -36,7 +36,7 @@ module.exports = class LlmErrorMessage {
     this.embedding_id = embedding?.id
     this.vector_store_id = vectorsearch?.id
     this.tool_id = tool?.id
-    this.agent_id = agent?.id
+    this.agent_id = aiAgent?.id
 
     if (embedding?.vendor === 'gemini' || summary?.vendor === 'gemini') {
       this._handleGemini(cause)

lib/subscribers/langgraph/base.js

@@ -3,32 +3,31 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-const LangGraphSubscriber = require('./base')
+const AiMonitoringSubscriber = require('../ai-monitoring/base')
 const { AI: { LANGGRAPH } } = require('#agentlib/metrics/names.js')
 
 /**
- * Subscribes to `Pregel.invoke()` events. Only exists for
- * a more accurate segment name. Its sibling `LangGraphStreamSubscriber`
+ * Subscribes to `Pregel.invoke()` events. Only exists to
+ * create a segment/span. The related `LangGraphStreamSubscriber`
  * handles the creation of Llm events because `invoke` is just a
  * wrapper around `stream`, and we don't want to create duplicates.
  */
-class LangGraphInvokeSubscriber extends LangGraphSubscriber {
+class LangGraphInvokeSubscriber extends AiMonitoringSubscriber {
   constructor({ agent, logger }) {
-    super({ agent, logger, channelName: 'nr_invoke' })
+    super({ agent,
+      logger,
+      packageName: '@langchain/langgraph',
+      channelName: 'nr_invoke',
+      name: `${LANGGRAPH.AGENT}/invoke`,
+      trackingPrefix: LANGGRAPH.TRACKING_PREFIX })
+    this.events = ['asyncEnd']
   }
 
   handler(data, ctx) {
-    if (!this.enabled) {
-      this.logger.debug('LangGraph instrumentation is disabled, not creating segment.')
-      return ctx
-    }
-
-    // Create segment and return the context with it.
-    const agentName = data?.self?.name ?? 'agent'
-    return this.createSegment({
-      name: `${LANGGRAPH.AGENT}/invoke/${agentName}`,
-      ctx
-    })
+    // Update segment name with LangGraph agent name.
+    const aiAgentName = data?.self?.name ?? 'agent'
+    this.name = `${LANGGRAPH.AGENT}/invoke/${aiAgentName}`
+    return super.handler(data, ctx)
   }
 }
 

lib/subscribers/langgraph/graph-invoke.js

@@ -3,33 +3,32 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-const LangGraphSubscriber = require('./base')
+const AiMonitoringSubscriber = require('../ai-monitoring/base')
 const { DESTINATIONS } = require('#agentlib/config/attribute-filter.js')
 const LlmErrorMessage = require('#agentlib/llm-events/error-message.js')
 const LangGraphAgentEvent = require('#agentlib/llm-events/langgraph/agent.js')
 const { AI: { LANGGRAPH } } = require('#agentlib/metrics/names.js')
 
-class LangGraphStreamSubscriber extends LangGraphSubscriber {
+class LangGraphStreamSubscriber extends AiMonitoringSubscriber {
   constructor({ agent, logger }) {
-    super({ agent, logger, channelName: 'nr_stream' })
+    super({ agent,
+      logger,
+      packageName: '@langchain/langgraph',
+      channelName: 'nr_stream',
+      name: `${LANGGRAPH.AGENT}/stream`,
+      trackingPrefix: LANGGRAPH.TRACKING_PREFIX })
+    this.events = ['asyncEnd']
   }
 
   get enabled() {
     return super.enabled && this.agent.config.ai_monitoring.streaming.enabled
   }
 
   handler(data, ctx) {
-    if (!this.enabled) {
-      this.logger.debug('LangGraph streaming instrumentation is disabled, not creating segment.')
-      return ctx
-    }
-
-    // Create segment and return the context with it.
-    const agentName = data?.self?.name ?? 'agent'
-    return this.createSegment({
-      name: `${LANGGRAPH.AGENT}/stream/${agentName}`,
-      ctx
-    })
+    // Update segment name with LangGraph agent name.
+    const aiAgentName = data?.self?.name ?? 'agent'
+    this.name = `${LANGGRAPH.AGENT}/stream/${aiAgentName}`
+    return super.handler(data, ctx)
   }
 
   asyncEnd(data) {
@@ -44,29 +43,31 @@ class LangGraphStreamSubscriber extends LangGraphSubscriber {
     if (transaction?.isActive() !== true) {
       return
     }
-    const { moduleVersion: pkgVersion, error: initialErr, result: stream } = data
-    const name = data?.self?.name ?? 'agent'
+
+    // Extract data
+    const { error: initialErr, result: stream } = data
+    const aiAgentName = data?.self?.name ?? 'agent'
 
     // Mark as LLM event
     transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
 
     // Add AI Agent component attribute
-    segment.addSpanAttribute('subcomponent', `{"type": "APM-AI_AGENT", "name": ${name}}`)
+    segment.addSpanAttribute('subcomponent', `{"type": "APM-AI_AGENT", "name": ${aiAgentName}}`)
 
     // If error already:
     if (initialErr) {
-      this.recordAgentEvent({ name, transaction, segment, pkgVersion, error: initialErr })
+      this.recordAiAgentEvent({ aiAgentName, transaction, segment, error: initialErr })
       return
     }
 
     // Note: as of 18.x `ReadableStream` is a global
     // eslint-disable-next-line n/no-unsupported-features/node-builtins
     if (stream instanceof ReadableStream) {
-      this.wrapNextHandler({ stream, segment, transaction, pkgVersion, name })
+      this.wrapNextHandler({ stream, segment, transaction, aiAgentName })
     }
   }
 
-  wrapNextHandler({ stream, segment, transaction, pkgVersion, name }) {
+  wrapNextHandler({ stream, segment, transaction, aiAgentName }) {
     const self = this
     const orig = stream.getReader
     stream.getReader = function wrappedGetReaderLangGraph() {
@@ -76,11 +77,11 @@ class LangGraphStreamSubscriber extends LangGraphSubscriber {
         try {
           const result = await origRead.apply(this, args)
           if (result?.done) {
-            self.recordAgentEvent({ name, transaction, segment, pkgVersion, error: false })
+            self.recordAiAgentEvent({ aiAgentName, transaction, segment, error: false })
           }
           return result
         } catch (err) {
-          self.recordAgentEvent({ name, transaction, segment, pkgVersion, error: err })
+          self.recordAiAgentEvent({ aiAgentName, transaction, segment, error: err })
           throw err
         } finally {
           // update segment duration on every stream iteration to extend
@@ -92,18 +93,28 @@ class LangGraphStreamSubscriber extends LangGraphSubscriber {
     }
   }
 
-  recordAgentEvent({ name, transaction, segment, pkgVersion, error }) {
+  /**
+   * Records a LangGraph `LlmAgent` event and a `LlmErrorMessage` event
+   * if an error occurred.
+   *
+   * @param {object} params function parameters
+   * @param {string} params.aiAgentName LangGraph AI agent name
+   * @param {object} params.transaction current transaction
+   * @param {object} params.segment current segment
+   * @param {object} [params.error] an error if it occurred
+   */
+  recordAiAgentEvent({ aiAgentName, transaction, segment, error }) {
     const { agent } = this
     segment.end()
     const agentEvent = new LangGraphAgentEvent({
       agent,
-      name,
+      name: aiAgentName,
       transaction,
       segment,
-      error: error !== null
+      error: !!error
     })
 
-    this.recordEvent({ type: 'LlmAgent', pkgVersion, msg: agentEvent })
+    this.recordEvent({ type: 'LlmAgent', msg: agentEvent })
 
     if (error) {
       agent.errors.add(
@@ -112,7 +123,7 @@ class LangGraphStreamSubscriber extends LangGraphSubscriber {
         new LlmErrorMessage({
           response: {},
           cause: error,
-          agent: agentEvent
+          aiAgent: agentEvent
         })
       )
     }