feat: Added instrumentation for langchain tools. (#2009)

lib/instrumentation/langchain/callback-manager.js

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { langchainRunId } = require('../../symbols')
+
+module.exports = function initialize(shim, callbacks) {
+  shim.wrap(
+    callbacks.CallbackManager.prototype,
+    ['handleChainStart', 'handleToolStart'],
+    function wrapStart(shim, orig) {
+      return async function wrappedStart() {
+        const result = await orig.apply(this, arguments)
+        const segment = shim.getActiveSegment()
+        if (segment) {
+          segment[langchainRunId] = result?.runId
+        }
+
+        return result
+      }
+    }
+  )
+}

lib/instrumentation/langchain/common.js

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const {
+  AI: { LANGCHAIN }
+} = require('../../metrics/names')
+
+const common = module.exports
+
+/**
+ * Langchain allows you to define tags at the instance and call level
+ * This helper merges the two into 1 array ensuring there are not duplicates
+ *
+ * @param {Array} localTags tags defined on instance of a langchain object
+ * @param {Array} paramsTags tags defined on the method
+ * @returns {Array} a merged array of unique tags
+ */
+common.mergeTags = function mergeTags(localTags = [], paramsTags = []) {
+  const tags = localTags.filter((tag) => !paramsTags.includes(tag))
+  tags.push(...paramsTags)
+  return tags
+}
+
+/**
+ * Langchain allows you to define metadata at the instance and call level
+ * This helper merges the two into object favoring the call level metadata
+ * values when duplicate keys exist.
+ *
+ * @param {object} localMeta metadata defined on instance of a langchain object
+ * @param {object} paramsMeta metadata defined on the method
+ * @returns {object} a merged object of metadata
+ */
+common.mergeMetadata = function mergeMetadata(localMeta = {}, paramsMeta = {}) {
+  return { ...localMeta, ...paramsMeta }
+}
+
+/**
+ * Helper to enqueue a LLM event into the custom event aggregator.  This will also
+ * increment the Supportability metric that's used to derive a tag on the APM entity.
+ *
+ * @param {object} params function params
+ * @param {Agent} params.agent NR agent
+ * @param {string} params.type type of llm event(i.e.- LlmChatCompletionMessage, LlmTool, etc)
+ * @param {object} params.msg the llm event getting enqueued
+ * @param {string} params.pkgVersion version of langchain library instrumented
+ */
+common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) {
+  agent.metrics.getOrCreateMetric(`${LANGCHAIN.TRACKING_PREFIX}/${pkgVersion}`).incrementCallCount()
+  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
+}
+
+/**
+ * Helper to decide if instrumentation should be registered.
+ *
+ * @param {object} config agent config
+ * @returns {boolean} flag if we should skip instrumentation
+ */
+common.shouldSkipInstrumentation = function shouldSkipInstrumentation(config) {
+  return !(
+    config.ai_monitoring.enabled === true && config.feature_flag.langchain_instrumentation === true
+  )
+}

lib/instrumentation/langchain/nr-hooks.js

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const toolsInstrumentation = require('./tools')
+const cbManagerInstrumentation = require('./callback-manager')
+
+module.exports = [
+  {
+    type: 'generic',
+    moduleName: '@langchain/core/tools',
+    onRequire: toolsInstrumentation
+  },
+  {
+    type: 'generic',
+    moduleName: '@langchain/core/dist/callbacks/manager',
+    onRequire: cbManagerInstrumentation
+  }
+]

lib/instrumentation/langchain/tools.js

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { LangChainTool } = require('../../llm-events/langchain')
+const { mergeMetadata, mergeTags, recordEvent, shouldSkipInstrumentation } = require('./common')
+const {
+  AI: { LANGCHAIN }
+} = require('../../metrics/names')
+const { langchainRunId } = require('../../symbols')
+const { DESTINATIONS } = require('../../config/attribute-filter')
+const { RecorderSpec } = require('../../shim/specs')
+
+module.exports = function initialize(shim, tools) {
+  const { agent, pkgVersion } = shim
+
+  if (shouldSkipInstrumentation(agent.config)) {
+    shim.logger.debug(
+      'langchain instrumentation is disabled.  To enable set `config.ai_monitoring.enabled` to true'
+    )
+    return
+  }
+
+  shim.record(tools.StructuredTool.prototype, 'call', function wrapCall(shim, call, fnName, args) {
+    const { name, metadata: instanceMeta, description, tags: instanceTags } = this
+    const [request, params] = args
+    const { metadata: paramsMeta, tags: paramsTags } = params || {}
+    return new RecorderSpec({
+      name: `${LANGCHAIN.TOOL}/${name}`,
+      promise: true,
+      // eslint-disable-next-line max-params
+      after(_shim, _fn, _name, _err, output, segment) {
+        const metadata = mergeMetadata(instanceMeta, paramsMeta)
+        const tags = mergeTags(instanceTags, paramsTags)
+        segment.end()
+        const toolEvent = new LangChainTool({
+          agent,
+          description,
+          name,
+          runId: segment[langchainRunId],
+          metadata,
+          tags,
+          input: request?.input,
+          output,
+          segment
+        })
+        recordEvent({ agent, type: 'LlmTool', pkgVersion, msg: toolEvent })
+        segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
+      }
+    })
+  })
+}

lib/instrumentations.js

@@ -25,11 +25,13 @@ module.exports = function instrumentations() {
     '@hapi/hapi': { type: MODULE_TYPE.WEB_FRAMEWORK },
     'ioredis': { type: MODULE_TYPE.DATASTORE },
     'koa': { module: '@newrelic/koa' },
+    'langchain': { module: './instrumentation/langchain' },
     'memcached': { type: MODULE_TYPE.DATASTORE },
     'mongodb': { type: MODULE_TYPE.DATASTORE },
     'mysql': { module: './instrumentation/mysql' },
     'openai': { type: MODULE_TYPE.GENERIC },
     '@nestjs/core': { type: MODULE_TYPE.WEB_FRAMEWORK },
+    '@prisma/client': { type: MODULE_TYPE.DATASTORE },
     'pino': { module: './instrumentation/pino' },
     'pg': { type: MODULE_TYPE.DATASTORE },
     'q': { type: null },
@@ -53,7 +55,6 @@ module.exports = function instrumentations() {
     'loglevel': { type: MODULE_TYPE.TRACKING },
     'npmlog': { type: MODULE_TYPE.TRACKING },
     'fancy-log': { type: MODULE_TYPE.TRACKING },
-    '@prisma/client': { type: MODULE_TYPE.DATASTORE },
     'knex': { type: MODULE_TYPE.TRACKING }
   }
 }

lib/llm-events/langchain/chat-completion-message.js

@@ -26,7 +26,7 @@ const defaultParams = {
   completionId: makeId(36)
 }
 
-class LangChainCompletionMesssage extends LangChainEvent {
+class LangChainCompletionMessage extends LangChainEvent {
   content
   role
   sequence
@@ -49,4 +49,4 @@ class LangChainCompletionMesssage extends LangChainEvent {
   }
 }
 
-module.exports = LangChainCompletionMesssage
+module.exports = LangChainCompletionMessage

lib/llm-events/langchain/chat-completion-summary.js

@@ -28,29 +28,14 @@ class LangChainCompletionSummary extends LangChainEvent {
   duration;
   ['response.number_of_messages'] = 0
 
-  #tags
-
   constructor(params = defaultParams) {
     params = Object.assign({}, defaultParams, params)
     super(params)
     const { segment } = params
 
-    this.tags = params.tags
     this.duration = segment?.getDurationInMillis()
     this['response.number_of_messages'] = params.messages?.length
   }
-
-  get tags() {
-    return this.#tags
-  }
-
-  set tags(value) {
-    if (Array.isArray(value)) {
-      this.#tags = value.join(',')
-    } else if (typeof value === 'string') {
-      this.#tags = value
-    }
-  }
 }
 
 module.exports = LangChainCompletionSummary

lib/llm-events/langchain/event.js

@@ -29,6 +29,7 @@ const defaultParams = {
   segment: {
     transaction: {}
   },
+  tags: [],
   runId: '',
   metadata: {},
   virtual: undefined
@@ -60,6 +61,7 @@ class LangChainEvent extends BaseEvent {
     this.trace_id = segment?.transaction?.traceId
     this.langchainMeta = params.metadata
     this.metadata = agent
+    this.tags = Array.isArray(params.tags) ? params.tags.join(',') : params.tags
 
     if (params.virtual !== undefined) {
       if (params.virtual !== true && params.virtual !== false) {

lib/llm-events/langchain/index.js

@@ -10,5 +10,6 @@ module.exports = {
   LangChainCompletionMessage: require('./chat-completion-message'),
   LangChainCompletionSummary: require('./chat-completion-summary'),
   LangChainVectorSearch: require('./vector-search'),
-  LangChainVectorSearchResult: require('./vector-search-result')
+  LangChainVectorSearchResult: require('./vector-search-result'),
+  LangChainTool: require('./tool')
 }

lib/llm-events/langchain/tool.js

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const LangChainEvent = require('./event')
+
+class LangChainTool extends LangChainEvent {
+  constructor(params) {
+    super(params)
+    this.input = params.input
+    this.output = params.output
+    this.name = params.name
+    this.description = params.description
+    this.duration = params?.segment?.getDurationInMillis()
+    this.run_id = this.request_id
+    delete this.request_id
+    delete this.virtual_llm
+    delete this.conversation_id
+  }
+}
+
+module.exports = LangChainTool

lib/metrics/names.js

@@ -168,7 +168,8 @@ const AI = {
   TRACKING_PREFIX: `${SUPPORTABILITY.NODEJS}/ML`,
   STREAMING_DISABLED: `${SUPPORTABILITY.NODEJS}/ML/Streaming/Disabled`,
   EMBEDDING: 'Llm/embedding',
-  COMPLETION: 'Llm/completion'
+  COMPLETION: 'Llm/completion',
+  TOOL: 'Llm/tool'
 }
 
 AI.OPENAI = {
@@ -177,6 +178,13 @@ AI.OPENAI = {
   COMPLETION: `${AI.COMPLETION}/OpenAI/create`
 }
 
+AI.LANGCHAIN = {
+  TRACKING_PREFIX: `${AI.TRACKING_PREFIX}/Langchain`,
+  EMBEDDING: `${AI.EMBEDDING}/Langchain`,
+  COMPLETION: `${AI.COMPLETION}/Langchain`,
+  TOOL: `${AI.TOOL}/Langchain`
+}
+
 const RESTIFY = {
   PREFIX: 'Restify/'
 }

lib/symbols.js

@@ -21,6 +21,7 @@ module.exports = {
   openAiHeaders: Symbol('openAiHeaders'),
   openAiApiKey: Symbol('openAiApiKey'),
   parentSegment: Symbol('parentSegment'),
+  langchainRunId: Symbol('runId'),
   prismaConnection: Symbol('prismaConnection'),
   prismaModelCall: Symbol('modelCall'),
   segment: Symbol('segment'),