chore(llmobs): add internal ability to tag tool id on tool messages (#6022)

* add ability to tag tool id on tool messages

* add type doc

* fix test

* remove documenting

Author: sabrenner

packages/dd-trace/src/llmobs/tagger.js

@@ -281,14 +281,15 @@ class LLMObsTagger {
 
       const { content = '', role } = message
       const toolCalls = message.toolCalls
+      const toolId = message.toolId
       const messageObj = { content }
 
       const valid = typeof content === 'string'
       if (!valid) {
         this.#handleFailure('Message content must be a string.', 'invalid_io_messages')
       }
 
-      const condition = this.#tagConditionalString(role, 'Message role', messageObj, 'role')
+      let condition = this.#tagConditionalString(role, 'Message role', messageObj, 'role')
 
       if (toolCalls) {
         const filteredToolCalls = this.#filterToolCalls(toolCalls)
@@ -298,6 +299,14 @@ class LLMObsTagger {
         }
       }
 
+      if (toolId) {
+        if (role === 'tool') {
+          condition = this.#tagConditionalString(toolId, 'Tool ID', messageObj, 'tool_id')
+        } else {
+          log.warn(`Tool ID for tool message not associated with a "tool" role, instead got "${role}"`)
+        }
+      }
+
       if (valid && condition) {
         messages.push(messageObj)
       }

packages/dd-trace/test/llmobs/tagger.spec.js

@@ -391,6 +391,44 @@ describe('tagger', () => {
           expect(() => tagger.tagLLMIO(span, messages, undefined)).to.throw()
         })
       })
+
+      describe('tool message tagging', () => {
+        it('tags a span with a tool message', () => {
+          const messages = [
+            { role: 'tool', content: 'The weather in San Francisco is sunny', toolId: '123' }
+          ]
+
+          tagger._register(span)
+          tagger.tagLLMIO(span, messages, undefined)
+          expect(Tagger.tagMap.get(span)).to.deep.equal({
+            '_ml_obs.meta.input.messages': [
+              { role: 'tool', content: 'The weather in San Francisco is sunny', tool_id: '123' }
+            ]
+          })
+        })
+
+        it('throws if the tool id is not a string', () => {
+          const messages = [
+            { role: 'tool', content: 'The weather in San Francisco is sunny', toolId: 123 }
+          ]
+
+          expect(() => tagger.tagLLMIO(span, messages, undefined)).to.throw()
+        })
+
+        it('logs a warning if the tool id is not associated with a tool role', () => {
+          const messages = [
+            { role: 'user', content: 'The weather in San Francisco is sunny', toolId: '123' }
+          ]
+
+          tagger._register(span)
+          tagger.tagLLMIO(span, messages, undefined)
+
+          const messageTags = Tagger.tagMap.get(span)['_ml_obs.meta.input.messages']
+          expect(messageTags[0]).to.not.have.property('tool_id')
+
+          expect(logger.warn).to.have.been.calledOnce
+        })
+      })
     })
 
     describe('tagEmbeddingIO', () => {
@@ -640,6 +678,17 @@ describe('tagger', () => {
           expect(logger.warn.callCount).to.equal(5) // 4 for tool call + 1 for role
         })
       })
+
+      it('logs a warning if the tool id is not a string', () => {
+        const messages = [
+          { role: 'tool', content: 'The weather in San Francisco is sunny', toolId: 123 }
+        ]
+
+        tagger._register(span)
+        tagger.tagLLMIO(span, messages, undefined)
+        expect(Tagger.tagMap.get(span)).to.not.have.property('_ml_obs.meta.input.messages')
+        expect(logger.warn).to.have.been.calledOnce
+      })
     })
   })
 })