feat(langchain, llmobs): expand langchain support for tools and vectorstores (#5760)

sabrenner · tlhunter · web-flow · commit 2244cf0e7642 · 2025-07-15T09:11:01.000-04:00
* minimal apm spans

* add llmobs plugins

* tool io refactor

* add tool tests

* apm vectorstore tests

* fix peer dependencies issue

* switch to agent.assertSomeTraces

* llmobs vectorstores tests

* remove .only

* remove promise.all

* lint

* do not need to format i/o for tools

---------

Co-authored-by: Thomas Hunter II &lt;tlhunter@datadog.com&gt;
diff --git a/packages/datadog-instrumentations/src/langchain.js b/packages/datadog-instrumentations/src/langchain.js
@@ -53,6 +53,27 @@ for (const extension of extensions) {
     return exports
   })
 
+  addHook({ name: '@langchain/core', file: `dist/tools/index.${extension}`, versions: ['>=0.1'] }, exports => {
+    if (extension === 'cjs') {
+      wrap(exports.StructuredTool.prototype, 'invoke', 'orchestrion:@langchain/core:Tool_invoke')
+    }
+    return exports
+  })
+
+  addHook({ name: '@langchain/core', file: `dist/vectorstores.${extension}`, versions: ['>=0.1'] }, exports => {
+    if (extension === 'cjs') {
+      wrap(
+        exports.VectorStore.prototype, 'similaritySearch', 'orchestrion:@langchain/core:VectorStore_similaritySearch'
+      )
+      wrap(
+        exports.VectorStore.prototype, 'similaritySearchWithScore',
+        'orchestrion:@langchain/core:VectorStore_similaritySearchWithScore'
+      )
+    }
+
+    return exports
+  })
+
   addHook({ name: '@langchain/core', file: `dist/embeddings.${extension}`, versions: ['>=0.1'] }, exports => {
     if (extension === 'cjs') {
       shimmer.wrap(exports, 'Embeddings', Embeddings => {
diff --git a/packages/datadog-instrumentations/src/orchestrion-config/index.js b/packages/datadog-instrumentations/src/orchestrion-config/index.js
@@ -53,4 +53,34 @@ instrumentations:
       class: Embeddings
     operator: traceSync
     channel_name: "Embeddings_constructor"
+  - module_name: "@langchain/core"
+    version_range: ">=0.1.0"
+    file_path: dist/tools/index.js
+    function_query:
+      name: invoke
+      type: method
+      kind: async
+      class: StructuredTool
+    operator: tracePromise
+    channel_name: "Tool_invoke"
+  - module_name: "@langchain/core"
+    version_range: ">=0.1.0"
+    file_path: dist/vectorstores.js
+    function_query:
+      name: similaritySearch
+      type: method
+      kind: async
+      class: VectorStore
+    operator: tracePromise
+    channel_name: "VectorStore_similaritySearch"
+  - module_name: "@langchain/core"
+    version_range: ">=0.1.0"
+    file_path: dist/vectorstores.js
+    function_query:
+      name: similaritySearchWithScore
+      type: method
+      kind: async
+      class: VectorStore
+    operator: tracePromise
+    channel_name: "VectorStore_similaritySearchWithScore"
 `
diff --git a/packages/datadog-plugin-langchain/src/tracing.js b/packages/datadog-plugin-langchain/src/tracing.js
@@ -40,9 +40,14 @@ class BaseLangChainTracingPlugin extends TracingPlugin {
 
     // Runnable interfaces have an `lc_namespace` property
     const ns = ctx.self.lc_namespace || ctx.namespace
-    const resource = ctx.resource = [...ns, ctx.self.constructor.name].join('.')
 
-    const handler = this.handlers[type]
+    const resourceParts = [...ns, ctx.self.constructor.name]
+    if (type === 'tool') {
+      resourceParts.push(ctx.instance.name)
+    }
+    const resource = ctx.resource = resourceParts.join('.')
+
+    const handler = this.handlers[type] || this.handlers.default
 
     const instance = ctx.instance
     const apiKey = handler.extractApiKey(instance)
@@ -78,7 +83,7 @@ class BaseLangChainTracingPlugin extends TracingPlugin {
 
     const { type } = ctx
 
-    const handler = this.handlers[type]
+    const handler = this.handlers[type] || this.handlers.default
     const tags = handler.getSpanEndTags(ctx, span) || {}
 
     span.addTags(tags)
@@ -139,11 +144,38 @@ class EmbeddingsEmbedDocumentsPlugin extends BaseLangChainTracingPlugin {
   }
 }
 
+class ToolInvokePlugin extends BaseLangChainTracingPlugin {
+  static get id () { return 'langchain_tool_invoke' }
+  static get lcType () { return 'tool' }
+  static get prefix () {
+    return 'tracing:orchestrion:@langchain/core:Tool_invoke'
+  }
+}
+
+class VectorStoreSimilaritySearchPlugin extends BaseLangChainTracingPlugin {
+  static get id () { return 'langchain_vectorstore_similarity_search' }
+  static get lcType () { return 'similarity_search' }
+  static get prefix () {
+    return 'tracing:orchestrion:@langchain/core:VectorStore_similaritySearch'
+  }
+}
+
+class VectorStoreSimilaritySearchWithScorePlugin extends BaseLangChainTracingPlugin {
+  static get id () { return 'langchain_vectorstore_similarity_search_with_score' }
+  static get lcType () { return 'similarity_search' }
+  static get prefix () {
+    return 'tracing:orchestrion:@langchain/core:VectorStore_similaritySearchWithScore'
+  }
+}
+
 module.exports = [
   RunnableSequenceInvokePlugin,
   RunnableSequenceBatchPlugin,
   BaseChatModelGeneratePlugin,
   BaseLLMGeneratePlugin,
   EmbeddingsEmbedQueryPlugin,
-  EmbeddingsEmbedDocumentsPlugin
+  EmbeddingsEmbedDocumentsPlugin,
+  ToolInvokePlugin,
+  VectorStoreSimilaritySearchPlugin,
+  VectorStoreSimilaritySearchWithScorePlugin
 ]
diff --git a/packages/datadog-plugin-langchain/test/index.spec.js b/packages/datadog-plugin-langchain/test/index.spec.js
@@ -19,6 +19,27 @@ const openAiBaseEmbeddingInfo = { base: 'https://api.openai.com', path: '/v1/emb
 
 const isDdTrace = iastFilter.isDdTrace
 
+function stubSingleEmbedding (langchainOpenaiOpenAiVersion) {
+  if (semver.satisfies(langchainOpenaiOpenAiVersion, '>=4.91.0')) {
+    stubCall({
+      ...openAiBaseEmbeddingInfo,
+      response: require('./fixtures/single-embedding.json')
+    })
+  } else {
+    stubCall({
+      ...openAiBaseEmbeddingInfo,
+      response: {
+        object: 'list',
+        data: [{
+          object: 'embedding',
+          index: 0,
+          embedding: Array(1536).fill(0)
+        }]
+      }
+    })
+  }
+}
+
 describe('Plugin', () => {
   let langchainOpenai
   let langchainAnthropic
@@ -28,7 +49,8 @@ describe('Plugin', () => {
   let langchainOutputParsers
   let langchainPrompts
   let langchainRunnables
-
+  let langchainTools
+  let MemoryVectorStore
   /**
    * In OpenAI 4.91.0, the default response format for embeddings was changed from `float` to `base64`.
    * We do not have control in @langchain/openai embeddings to change this for an individual call,
@@ -65,7 +87,8 @@ describe('Plugin', () => {
       })
 
       beforeEach(() => {
-        langchainOpenai = require(`../../../versions/@langchain/openai@${version}`).get()
+        langchainOpenai = require(`../../../versions/langchain@${version}`)
+          .get('@langchain/openai')
         langchainAnthropic = require(`../../../versions/@langchain/anthropic@${version}`).get()
         if (version !== '0.1.0') {
           // version mismatching otherwise
@@ -80,10 +103,17 @@ describe('Plugin', () => {
         langchainPrompts = require(`../../../versions/@langchain/core@${version}`).get('@langchain/core/prompts')
         langchainRunnables = require(`../../../versions/@langchain/core@${version}`).get('@langchain/core/runnables')
 
+        langchainTools = require(`../../../versions/@langchain/core@${version}`)
+          .get('@langchain/core/tools')
+
+        MemoryVectorStore = require(`../../../versions/@langchain/core@${version}`)
+          .get('langchain/vectorstores/memory')
+          .MemoryVectorStore
+
         langchainOpenaiOpenAiVersion =
-            require(`../../../versions/@langchain/openai@${version}`)
-              .get('openai/version')
-              .VERSION
+          require(`../../../versions/langchain@${version}`)
+            .get('openai/version')
+            .VERSION
       })
 
       afterEach(() => {
@@ -1013,6 +1043,133 @@ describe('Plugin', () => {
           })
         })
       })
+
+      describe('tools', () => {
+        it('traces a tool call', async function () {
+          if (!langchainTools?.tool) this.skip()
+
+          const myTool = langchainTools.tool(
+            () => 'Hello, world!',
+            {
+              name: 'myTool',
+              description: 'A tool that returns a greeting'
+            }
+          )
+
+          const checkTraces = agent.assertSomeTraces(traces => {
+            const span = traces[0][0]
+
+            expect(span).to.have.property('name', 'langchain.request')
+            expect(span.resource).to.match(/^langchain\.tools\.[^.]+\.myTool$/)
+          })
+          const result = await myTool.invoke()
+          expect(result).to.equal('Hello, world!')
+
+          await checkTraces
+        })
+
+        it('traces a tool call with an error', async function () {
+          if (!langchainTools?.tool) this.skip()
+
+          const myTool = langchainTools.tool(
+            () => { throw new Error('This is a test error') },
+            {
+              name: 'myTool',
+              description: 'A tool that throws an error'
+            }
+          )
+
+          const checkTraces = agent.assertSomeTraces(traces => {
+            const span = traces[0][0]
+
+            expect(span).to.have.property('name', 'langchain.request')
+            expect(span.resource).to.match(/^langchain\.tools\.[^.]+\.myTool$/)
+
+            expect(span.meta).to.have.property('error.message')
+            expect(span.meta).to.have.property('error.type')
+            expect(span.meta).to.have.property('error.stack')
+          })
+
+          try {
+            await myTool.invoke()
+            expect.fail('Expected an error to be thrown')
+          } catch {}
+
+          await checkTraces
+        })
+      })
+
+      describe('vectorstores', () => {
+        let vectorstore
+
+        beforeEach(async () => {
+          // need to mock out adding a document to the vectorstore
+          stubSingleEmbedding(langchainOpenaiOpenAiVersion)
+
+          const embeddings = new langchainOpenai.OpenAIEmbeddings()
+          vectorstore = new MemoryVectorStore(embeddings)
+
+          const document = {
+            pageContent: 'The powerhouse of the cell is the mitochondria',
+            metadata: { source: 'https://example.com' },
+            id: '1'
+          }
+
+          return vectorstore.addDocuments([document])
+        })
+
+        it('traces a vectorstore similaritySearch call', async () => {
+          stubSingleEmbedding(langchainOpenaiOpenAiVersion)
+
+          const checkTraces = agent.assertSomeTraces(traces => {
+            const spans = traces[0]
+
+            expect(spans).to.have.length(2)
+
+            const vectorstoreSpan = spans[0]
+            const embeddingSpan = spans[1]
+
+            expect(vectorstoreSpan).to.have.property('name', 'langchain.request')
+            expect(vectorstoreSpan).to.have.property('resource', 'langchain.vectorstores.memory.MemoryVectorStore')
+
+            expect(embeddingSpan).to.have.property('name', 'langchain.request')
+            expect(embeddingSpan).to.have.property('resource', 'langchain.embeddings.openai.OpenAIEmbeddings')
+          }, { spanResourceMatch: /langchain\.vectorstores\.memory\.MemoryVectorStore/ })
+          // we need the spanResourceMatch, otherwise we'll match from the beforeEach
+
+          const result = await vectorstore.similaritySearch('The powerhouse of the cell is the mitochondria', 2)
+          expect(result).to.exist
+
+          await checkTraces
+        })
+
+        it('traces a vectorstore similaritySearchWithScore call', async () => {
+          stubSingleEmbedding(langchainOpenaiOpenAiVersion)
+
+          const checkTraces = agent.assertSomeTraces(traces => {
+            const spans = traces[0]
+
+            expect(spans).to.have.length(2)
+
+            const vectorstoreSpan = spans[0]
+            const embeddingSpan = spans[1]
+
+            expect(vectorstoreSpan).to.have.property('name', 'langchain.request')
+            expect(vectorstoreSpan).to.have.property('resource', 'langchain.vectorstores.memory.MemoryVectorStore')
+
+            expect(embeddingSpan).to.have.property('name', 'langchain.request')
+            expect(embeddingSpan).to.have.property('resource', 'langchain.embeddings.openai.OpenAIEmbeddings')
+          }, { spanResourceMatch: /langchain\.vectorstores\.memory\.MemoryVectorStore/ })
+          // we need the spanResourceMatch, otherwise we'll match from the beforeEach
+
+          const result = await vectorstore.similaritySearchWithScore(
+            'The powerhouse of the cell is the mitochondria', 2
+          )
+          expect(result).to.exist
+
+          await checkTraces
+        })
+      })
     })
   })
 })
diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js
@@ -8,9 +8,14 @@ const ROLE_MAPPINGS = {
 
 class LangChainLLMObsHandler {
   constructor (tagger) {
+    /** @type {import('../../../tagger')} */
     this._tagger = tagger
   }
 
+  getName ({ span }) {
+    return span?.context()._tags?.['resource.name']
+  }
+
   setMetaTags () {}
 
   formatIO (messages) {
diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/tool.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/tool.js
@@ -0,0 +1,15 @@
+'use strict'
+
+const LangChainLLMObsHandler = require('.')
+
+class LangChainLLMObsToolHandler extends LangChainLLMObsHandler {
+  getName ({ instance }) {
+    return instance.name
+  }
+
+  setMetaTags ({ span, inputs, results }) {
+    this._tagger.tagTextIO(span, inputs, results)
+  }
+}
+
+module.exports = LangChainLLMObsToolHandler
diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/vectorstore.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/vectorstore.js
@@ -0,0 +1,36 @@
+'use strict'
+
+const LangChainLLMObsHandler = require('.')
+const { spanHasError } = require('../../../util')
+
+class LangChainLLMObsVectorStoreHandler extends LangChainLLMObsHandler {
+  setMetaTags ({ span, inputs, results }) {
+    const input = this.formatIO(inputs)
+    if (spanHasError(span)) {
+      this._tagger.tagRetrievalIO(span, input)
+      return
+    }
+
+    const documents = []
+    for (const documentResult of results) {
+      let document, score
+      if (Array.isArray(documentResult)) {
+        document = documentResult[0]
+        score = documentResult[1]
+      } else {
+        document = documentResult
+      }
+
+      documents.push({
+        text: document.pageContent,
+        id: document.id,
+        name: document.metadata?.source,
+        score
+      })
+    }
+
+    this._tagger.tagRetrievalIO(span, input, documents)
+  }
+}
+
+module.exports = LangChainLLMObsVectorStoreHandler
diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/index.js b/packages/dd-trace/src/llmobs/plugins/langchain/index.js
diff --git a/packages/dd-trace/test/llmobs/plugins/langchain/index.spec.js b/packages/dd-trace/test/llmobs/plugins/langchain/index.spec.js
diff --git a/packages/dd-trace/test/plugins/externals.json b/packages/dd-trace/test/plugins/externals.json
