fix(llmobs): render BaseMessage in langgraph workflow I/O (#8097)

* fix(llmobs): render BaseMessage in langgraph workflow I/O (#8096)

Extract the langchain handler's formatIO/getContentFromMessage/getRole into a shared util under and use it for both plugins, so workflow spans and child chat_model spans render messages consistently as { content, role }.

Author: crysmags

packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js

@@ -1,16 +1,17 @@
 'use strict'
 
 const { spanHasError } = require('../../../util')
+const { formatIO } = require('../messages')
 const LangChainLLMObsHandler = require('.')
 
 class LangChainLLMObsChainHandler extends LangChainLLMObsHandler {
   setMetaTags ({ span, inputs, results }) {
     let input
     if (inputs) {
-      input = this.formatIO(inputs)
+      input = formatIO(inputs)
     }
 
-    const output = !results || spanHasError(span) ? '' : this.formatIO(results)
+    const output = !results || spanHasError(span) ? '' : formatIO(results)
 
     // chain spans will always be workflows
     this._tagger.tagTextIO(span, input, output)

packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js

@@ -2,6 +2,7 @@
 
 const LLMObsTagger = require('../../../tagger')
 const { spanHasError } = require('../../../util')
+const { getRole } = require('../messages')
 const LangChainLLMObsHandler = require('.')
 
 const LLM = 'llm'
@@ -22,7 +23,7 @@ class LangChainLLMObsChatModelHandler extends LangChainLLMObsHandler {
     for (const messageSet of inputs) {
       for (const message of messageSet) {
         const content = message.content || ''
-        const role = this.getRole(message)
+        const role = getRole(message)
         inputMessages.push({ content, role })
       }
     }
@@ -54,7 +55,7 @@ class LangChainLLMObsChatModelHandler extends LangChainLLMObsHandler {
     for (const messageSet of results.generations) {
       for (const chatCompletion of messageSet) {
         const chatCompletionMessage = chatCompletion.message
-        const role = this.getRole(chatCompletionMessage)
+        const role = getRole(chatCompletionMessage)
         const content = chatCompletionMessage.text || ''
         const toolCalls = this.extractToolCalls(chatCompletionMessage)
         outputMessages.push({ content, role, toolCalls })

packages/dd-trace/src/llmobs/plugins/langchain/handlers/embedding.js

@@ -2,6 +2,7 @@
 
 const LLMObsTagger = require('../../../tagger')
 const { spanHasError } = require('../../../util')
+const { formatIO } = require('../messages')
 const LangChainLLMObsHandler = require('.')
 
 class LangChainLLMObsEmbeddingHandler extends LangChainLLMObsHandler {
@@ -10,7 +11,7 @@ class LangChainLLMObsEmbeddingHandler extends LangChainLLMObsHandler {
     let embeddingInput, embeddingOutput
 
     if (isWorkflow) {
-      embeddingInput = this.formatIO(inputs)
+      embeddingInput = formatIO(inputs)
     } else {
       const input = Array.isArray(inputs) ? inputs : [inputs]
       embeddingInput = input.map(doc => ({ text: doc }))

packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js

@@ -1,11 +1,5 @@
 'use strict'
 
-const ROLE_MAPPINGS = {
-  human: 'user',
-  ai: 'assistant',
-  system: 'system',
-}
-
 class LangChainLLMObsHandler {
   constructor (tagger) {
     /** @type {import('../../../tagger')} */
@@ -18,38 +12,6 @@ class LangChainLLMObsHandler {
 
   setMetaTags () {}
 
-  formatIO (messages) {
-    if (messages.constructor.name === 'Object') { // plain JSON
-      const formatted = {}
-      for (const [key, value] of Object.entries(messages)) {
-        formatted[key] = this.formatIO(value)
-      }
-
-      return formatted
-    } else if (Array.isArray(messages)) {
-      return messages.map(message => this.formatIO(message))
-    } // either a BaseMesage type or a string
-    return this.getContentFromMessage(messages)
-  }
-
-  getContentFromMessage (message) {
-    if (typeof message === 'string') {
-      return message
-    }
-    try {
-      const messageContent = {
-        content: message.content || '',
-      }
-
-      const role = this.getRole(message)
-      if (role) messageContent.role = role
-
-      return messageContent
-    } catch {
-      return JSON.stringify(message)
-    }
-  }
-
   checkTokenUsageChatOrLLMResult (results) {
     const llmOutput = results.llmOutput
     const tokens = {
@@ -90,17 +52,6 @@ class LangChainLLMObsHandler {
       runId: runIdBase,
     }
   }
-
-  getRole (message) {
-    if (message.role) return ROLE_MAPPINGS[message.role] || message.role
-
-    const type = (
-      (typeof message.getType === 'function' && message.getType()) ||
-      (typeof message._getType === 'function' && message._getType())
-    )
-
-    return ROLE_MAPPINGS[type] || type
-  }
 }
 
 module.exports = LangChainLLMObsHandler

packages/dd-trace/src/llmobs/plugins/langchain/handlers/vectorstore.js

@@ -1,11 +1,12 @@
 'use strict'
 
 const { spanHasError } = require('../../../util')
+const { formatIO } = require('../messages')
 const LangChainLLMObsHandler = require('.')
 
 class LangChainLLMObsVectorStoreHandler extends LangChainLLMObsHandler {
   setMetaTags ({ span, inputs, results }) {
-    const input = this.formatIO(inputs)
+    const input = formatIO(inputs)
     if (spanHasError(span)) {
       this._tagger.tagRetrievalIO(span, input)
       return

packages/dd-trace/src/llmobs/plugins/langchain/messages.js

@@ -0,0 +1,76 @@
+'use strict'
+
+const ROLE_MAPPINGS = {
+  human: 'user',
+  ai: 'assistant',
+  system: 'system',
+}
+
+function getRole (message) {
+  if (message.role) return ROLE_MAPPINGS[message.role] || message.role
+
+  const type = (
+    (typeof message.getType === 'function' && message.getType()) ||
+    (typeof message._getType === 'function' && message._getType())
+  )
+
+  return ROLE_MAPPINGS[type] || type
+}
+
+function getContentFromMessage (message) {
+  if (typeof message === 'string') {
+    return message
+  }
+  try {
+    const messageContent = {
+      content: message.content || '',
+    }
+
+    const role = getRole(message)
+    if (role) messageContent.role = role
+
+    return messageContent
+  } catch {
+    return JSON.stringify(message)
+  }
+}
+
+function isBaseMessage (data) {
+  return typeof data._getType === 'function' || typeof data.getType === 'function'
+}
+
+function formatIO (data) {
+  if (data == null) return ''
+
+  if (typeof data === 'string' || typeof data === 'number' || typeof data === 'boolean') {
+    return data
+  }
+
+  if (data.constructor?.name === 'Object') {
+    const formatted = {}
+    for (const [key, value] of Object.entries(data)) {
+      formatted[key] = formatIO(value)
+    }
+    return formatted
+  }
+
+  if (Array.isArray(data)) {
+    return data.map(item => formatIO(item))
+  }
+
+  // Only duck-typed BaseMessage instances collapse to { content, role }.
+  // Other class instances (e.g. LangChain Document) preserve their shape via JSON.stringify,
+  // otherwise they'd reduce to { content: '' } and lose data.
+  if (isBaseMessage(data)) return getContentFromMessage(data)
+
+  try {
+    return JSON.stringify(data)
+  } catch {
+    return String(data)
+  }
+}
+
+module.exports = {
+  getRole,
+  formatIO,
+}

packages/dd-trace/src/llmobs/plugins/langgraph/index.js

@@ -1,36 +1,11 @@
 'use strict'
 
 const LLMObsPlugin = require('../base')
+const { formatIO } = require('../langchain/messages')
 const { spanHasError } = require('../../util')
 
 const streamDataMap = new WeakMap()
 
-function formatIO (data) {
-  if (data == null) return ''
-
-  if (typeof data === 'string' || typeof data === 'number' || typeof data === 'boolean') {
-    return data
-  }
-
-  if (data.constructor?.name === 'Object') {
-    const formatted = {}
-    for (const [key, value] of Object.entries(data)) {
-      formatted[key] = formatIO(value)
-    }
-    return formatted
-  }
-
-  if (Array.isArray(data)) {
-    return data.map(item => formatIO(item))
-  }
-
-  try {
-    return JSON.stringify(data)
-  } catch {
-    return String(data)
-  }
-}
-
 class PregelStreamLLMObsPlugin extends LLMObsPlugin {
   static id = 'llmobs_langgraph_pregel_stream'
   static integration = 'langgraph'

packages/dd-trace/test/llmobs/plugins/langgraph/index.spec.js

@@ -13,6 +13,7 @@ const {
 describe('integrations', () => {
   let StateGraph
   let Annotation
+  let langchainMessages
 
   describe('langgraph', () => {
     const { getEvents } = useLlmObs({ plugin: ['langgraph', 'langchain'] })
@@ -22,6 +23,8 @@ describe('integrations', () => {
         const langgraph = require(`../../../../../../versions/@langchain/langgraph@${version}`).get()
         StateGraph = langgraph.StateGraph
         Annotation = langgraph.Annotation
+        langchainMessages = require(`../../../../../../versions/@langchain/langgraph@${version}`)
+          .get('@langchain/core/messages')
       })
 
       describe('Pregel.stream', () => {
@@ -152,6 +155,54 @@ describe('integrations', () => {
           })
         })
 
+        // Regression for https://github.com/DataDog/dd-trace-js/issues/8096: BaseMessage
+        // instances must render as { content, role } instead of full class dumps.
+        it('renders BaseMessage input/output as clean { content, role }', async () => {
+          const StateAnnotation = Annotation.Root({
+            messages: Annotation({
+              reducer: (x, y) => x.concat(y),
+              default: () => [],
+            }),
+          })
+
+          function chatNode () {
+            return {
+              messages: [new langchainMessages.AIMessage('Pong')],
+            }
+          }
+
+          const workflow = new StateGraph(StateAnnotation)
+            .addNode('chat', chatNode)
+            .addEdge('__start__', 'chat')
+            .addEdge('chat', '__end__')
+
+          const app = workflow.compile({ name: 'basemessage-graph' })
+
+          const chunks = []
+          for await (const chunk of await app.stream({
+            messages: [new langchainMessages.HumanMessage('Ping')],
+          })) {
+            chunks.push(chunk)
+          }
+
+          assert.ok(chunks.length > 0)
+
+          const { llmobsSpans } = await getEvents()
+
+          const workflowSpan = llmobsSpans.find(s => s.name === 'basemessage-graph')
+          assert.ok(workflowSpan, 'expected workflow span named basemessage-graph')
+
+          assert.strictEqual(
+            workflowSpan.meta.input.value,
+            JSON.stringify({ messages: [{ content: 'Ping', role: 'user' }] })
+          )
+
+          const parsedOutput = JSON.parse(workflowSpan.meta.output.value)
+          assert.ok(Array.isArray(parsedOutput.messages))
+          const lastMessage = parsedOutput.messages[parsedOutput.messages.length - 1]
+          assert.deepStrictEqual(lastMessage, { content: 'Pong', role: 'assistant' })
+        })
+
         it('does not tag output if streaming encounters an error', async () => {
           const StateAnnotation = Annotation.Root({
             value: Annotation({ default: () => 0 }),