feat: Add message entities for LangChain (#1983)

lib/llm-events/event.js

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const { DESTINATIONS } = require('../config/attribute-filter')
+
+class BaseLlmEvent {
+  conversationId(agent) {
+    const transaction = agent.tracer.getTransaction()
+    const attrs = transaction?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE)
+    return attrs?.['llm.conversation_id']
+  }
+}
+
+module.exports = BaseLlmEvent

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const LangChainEvent = require('./event')
+const { makeId } = require('../../util/hashes')
+
+/**
+ * @typedef {object} LangChainCompletionMessageParams
+ * @augments LangChainEventParams
+ * @property {string} content The text of the response received from LangChain.
+ * @property {string} role The role of the message, e.g. "human."
+ * @property {number} [sequence=0] The order of the message in the response.
+ * @property {string} [completionId] An identifier for the message.
+ */
+/**
+ * @type {LangChainCompletionMessageParams}
+ */
+const defaultParams = {
+  content: '',
+  role: undefined,
+  sequence: 0,
+  completionId: makeId(36)
+}
+
+class LangChainCompletionMesssage extends LangChainEvent {
+  content
+  role
+  sequence
+  completion_id
+
+  constructor(params = defaultParams) {
+    params = Object.assign({}, defaultParams, params)
+    super(params)
+
+    if (params.runId) {
+      this.id = `${params.runId}-${params.sequence}`
+    } else {
+      this.id = `${this.id}-${params.sequence}`
+    }
+
+    this.content = params.content
+    this.role = params.role
+    this.sequence = params.sequence
+    this.completion_id = params.completionId
+  }
+}
+
+module.exports = LangChainCompletionMesssage

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const LangChainEvent = require('./event')
+
+/**
+ * @typedef {object} LangChainCompletionSummaryParams
+ * @augments LangChainEventParams
+ * @property {string[]|string} [tags] A set of tags applied to the LangChain
+ * event. If provided as a simple string, it should be a comma separated value
+ * string.
+ * @property {object[]} messages The set of messages that were returned as the
+ * LangChain result.
+ */
+/**
+ * @type {LangChainCompletionSummaryParams}
+ */
+const defaultParams = {
+  tags: [],
+  messages: []
+}
+
+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

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const BaseEvent = require('../event')
+const { makeId } = require('../../util/hashes')
+const { isSimpleObject } = require('../../util/objects')
+
+/**
+ * @typedef {object} LangChainEventParams
+ * @property {object} agent A New Relic agent instance.
+ * @property {object} segment A New Relic segment instance.
+ * @property {string} runId The identifier LangChain has assigned to the "run."
+ * @property {Object<string, string>} metadata The metadata, if any, associated with the
+ * LangChain run.
+ * @property {boolean|undefined} [virtual=true] Indicates that this event is a
+ * LangChain specific event (`true`). LangChain is not itself an LLM, but an
+ * interface to many LLMs. Any LLMs LangChain interacts with that we have
+ * instrumented will have their own traces that are not "virtual."
+ */
+/**
+ * @type {LangChainEventParams}
+ */
+const defaultParams = {
+  agent: {},
+  segment: {
+    transaction: {}
+  },
+  runId: '',
+  metadata: {},
+  virtual: undefined
+}
+
+/**
+ * Baseline object representing an event in a LangChain conversation.
+ */
+class LangChainEvent extends BaseEvent {
+  id = makeId(36)
+  appName
+  conversation_id
+  span_id
+  request_id
+  transaction_id
+  trace_id
+  ingest_source = 'Node'
+  vendor = 'langchain'
+  virtual_llm = true
+
+  constructor(params = defaultParams) {
+    params = Object.assign({}, defaultParams, params)
+    super(params)
+    const { agent, segment } = params
+
+    this.appName = agent.config.applications()[0]
+    this.conversation_id = this.conversationId(agent)
+    this.span_id = segment?.id
+    this.request_id = params.runId
+    this.transaction_id = segment?.transaction?.id
+    this.trace_id = segment?.transaction?.traceId
+    this.metadata = params.metadata
+
+    if (params.virtual !== undefined) {
+      if (params.virtual !== true && params.virtual !== false) {
+        throw Error('params.virtual must be a primitive boolean')
+      }
+      this.virtual_llm = params.virtual
+    }
+  }
+
+  set metadata(value) {
+    if (isSimpleObject(value) === false) {
+      return
+    }
+    for (const [key, val] of Object.entries(value)) {
+      this[`metadata.${key}`] = val
+    }
+  }
+}
+
+module.exports = LangChainEvent

lib/llm-events/langchain/index.js

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+module.exports = {
+  LangChainEvent: require('./event'),
+  LangChainCompletionMessage: require('./chat-completion-message'),
+  LangChainCompletionSummary: require('./chat-completion-summary')
+}

lib/llm-events/openai/event.js

@@ -4,11 +4,14 @@
  */
 
 'use strict'
+
+const BaseEvent = require('../event')
 const { makeId } = require('../../util/hashes')
-const { DESTINATIONS } = require('../../../lib/config/attribute-filter')
 
-module.exports = class LlmEvent {
+module.exports = class LlmEvent extends BaseEvent {
   constructor({ agent, segment, request, response, responseAttrs = false }) {
+    super()
+
     this.id = makeId(36)
     this.appName = agent.config.applications()[0]
     this.request_id = response?.headers?.['x-request-id']
@@ -47,10 +50,4 @@ module.exports = class LlmEvent {
     this['response.headers.ratelimitRemainingRequests'] =
       response?.headers?.['x-ratelimit-remaining-requests']
   }
-
-  conversationId(agent) {
-    const transaction = agent.tracer.getTransaction()
-    const attrs = transaction?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE)
-    return attrs?.['llm.conversation_id']
-  }
 }

test/unit/llm-events/langchain/chat-completion-message.test.js

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const tap = require('tap')
+const LangChainCompletionMessage = require('../../../../lib/llm-events/langchain/chat-completion-message')
+
+tap.beforeEach((t) => {
+  t.context._tx = {
+    trace: {
+      custom: {
+        get() {
+          return {
+            'llm.conversation_id': 'test-conversation'
+          }
+        }
+      }
+    }
+  }
+
+  t.context.agent = {
+    config: {
+      applications() {
+        return ['test-app']
+      }
+    },
+    tracer: {
+      getTransaction() {
+        return t.context._tx
+      }
+    }
+  }
+
+  t.context.segment = {
+    id: 'segment-1',
+    transaction: {
+      id: 'tx-1',
+      traceId: 'trace-1'
+    }
+  }
+
+  t.context.runId = 'run-1'
+  t.context.metadata = { foo: 'foo' }
+})
+
+tap.test('creates entity', async (t) => {
+  const msg = new LangChainCompletionMessage({
+    ...t.context,
+    role: 'human',
+    sequence: 1,
+    content: 'hello world'
+  })
+  t.match(msg, {
+    id: 'run-1-1',
+    appName: 'test-app',
+    conversation_id: 'test-conversation',
+    span_id: 'segment-1',
+    request_id: 'run-1',
+    transaction_id: 'tx-1',
+    trace_id: 'trace-1',
+    ['metadata.foo']: 'foo',
+    ingest_source: 'Node',
+    vendor: 'langchain',
+    virtual_llm: true,
+    role: 'human',
+    sequence: 1,
+    content: 'hello world',
+    completion_id: /[a-z0-9-]{36}/
+  })
+})
+
+tap.test('assigns id correctly', async (t) => {
+  let msg = new LangChainCompletionMessage({ ...t.context, runId: '', sequence: 1 })
+  t.match(msg.id, /[a-z0-9-]{36}-1/)
+
+  msg = new LangChainCompletionMessage({ ...t.context, runId: '123456', sequence: 42 })
+  t.equal(msg.id, '123456-42')
+})