Skip to content

Commit f84216a

Browse files
Yun-KimYun-Kim
authored
feat(llmobs): add integration tag to llmobs spans (#5465)
* Add integration tag to llmobs spans * Add integration tag to tests
1 parent 579d213 commit f84216a

File tree

8 files changed

+51
-44
lines changed

8 files changed

+51
-44
lines changed

‎packages/dd-trace/src/llmobs/plugins/base.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ class LLMObsPlugin extends TracingPlugin {
4343
llmobsStorage.enterWith({ span })
4444
ctx.llmobs.parent = parent
4545

46-
this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions })
46+
this._tagger.registerLLMObsSpan(span, { parent, integration: this.constructor.id, ...registerOptions })
4747
}
4848
}
4949

‎packages/dd-trace/src/llmobs/plugins/bedrockruntime.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,8 @@ class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
5555
modelName: modelName.toLowerCase(),
5656
modelProvider: modelProvider.toLowerCase(),
5757
kind: 'llm',
58-
name: 'bedrock-runtime.command'
58+
name: 'bedrock-runtime.command',
59+
integration: 'bedrock'
5960
})
6061

6162
const requestParams = extractRequestParams(request.params, modelProvider)

‎packages/dd-trace/src/llmobs/span_processor.js‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ const {
77
METADATA,
88
INPUT_MESSAGES,
99
INPUT_VALUE,
10+
INTEGRATION,
1011
OUTPUT_MESSAGES,
1112
INPUT_DOCUMENTS,
1213
OUTPUT_DOCUMENTS,
@@ -186,6 +187,8 @@ class LLMObsSpanProcessor {
186187
const errType = span.context()._tags[ERROR_TYPE] || error?.name
187188
if (errType) tags.error_type = errType
188189
if (sessionId) tags.session_id = sessionId
190+
const integration = LLMObsTagger.tagMap.get(span)?.[INTEGRATION]
191+
if (integration) tags.integration = integration
189192
const existingTags = LLMObsTagger.tagMap.get(span)?.[TAGS] || {}
190193
if (existingTags) tags = { ...tags, ...existingTags }
191194
return Object.entries(tags).map(([key, value]) => `${key}:${value ?? ''}`)

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ const {
1010
INPUT_VALUE,
1111
OUTPUT_DOCUMENTS,
1212
INPUT_DOCUMENTS,
13+
INTEGRATION,
1314
OUTPUT_VALUE,
1415
METADATA,
1516
METRICS,
@@ -51,7 +52,8 @@ class LLMObsTagger {
5152
mlApp,
5253
parent,
5354
kind,
54-
name
55+
name,
56+
integration
5557
} = {}) {
5658
if (!this._config.llmobs.enabled) return
5759
if (!kind) return // do not register it in the map if it doesn't have an llmobs span kind
@@ -66,6 +68,7 @@ class LLMObsTagger {
6668

6769
sessionId = sessionId || registry.get(parent)?.[SESSION_ID]
6870
if (sessionId) this._setTag(span, SESSION_ID, sessionId)
71+
if (integration) this._setTag(span, INTEGRATION, integration)
6972

7073
if (!mlApp) mlApp = registry.get(parent)?.[ML_APP] || this._config.llmobs.mlApp
7174
this._setTag(span, ML_APP, mlApp)

‎packages/dd-trace/test/llmobs/plugins/aws-sdk/bedrockruntime.spec.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ describe('Plugin', () => {
109109
temperature: modelConfig.temperature,
110110
max_tokens: modelConfig.maxTokens
111111
},
112-
tags: { ml_app: 'test', language: 'javascript' }
112+
tags: { ml_app: 'test', language: 'javascript', integration: 'bedrock' }
113113
})
114114

115115
expect(spanEvent).to.deepEqualWithMockValues(expected)

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

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ describe('integrations', () => {
123123
outputMessages: [{ content: 'Hello, world!' }],
124124
metadata: MOCK_ANY,
125125
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
126-
tags: { ml_app: 'test', language: 'javascript' }
126+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
127127
})
128128

129129
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -151,7 +151,7 @@ describe('integrations', () => {
151151
outputMessages: [{ content: '' }],
152152
metadata: MOCK_ANY,
153153
tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 },
154-
tags: { ml_app: 'test', language: 'javascript' },
154+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' },
155155
error: 1,
156156
errorType: 'Error',
157157
errorMessage: MOCK_STRING,
@@ -210,7 +210,7 @@ describe('integrations', () => {
210210
metadata: MOCK_ANY,
211211
// @langchain/cohere does not provide token usage in the response
212212
tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 },
213-
tags: { ml_app: 'test', language: 'javascript' }
213+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
214214
})
215215

216216
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -255,7 +255,7 @@ describe('integrations', () => {
255255
outputMessages: [{ content: 'Hello, world!', role: 'assistant' }],
256256
metadata: MOCK_ANY,
257257
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
258-
tags: { ml_app: 'test', language: 'javascript' }
258+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
259259
})
260260

261261
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -283,7 +283,7 @@ describe('integrations', () => {
283283
outputMessages: [{ content: '' }],
284284
metadata: MOCK_ANY,
285285
tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 },
286-
tags: { ml_app: 'test', language: 'javascript' },
286+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' },
287287
error: 1,
288288
errorType: 'Error',
289289
errorMessage: MOCK_STRING,
@@ -334,7 +334,7 @@ describe('integrations', () => {
334334
outputMessages: [{ content: 'Hello!', role: 'assistant' }],
335335
metadata: MOCK_ANY,
336336
tokenMetrics: { input_tokens: 11, output_tokens: 6, total_tokens: 17 },
337-
tags: { ml_app: 'test', language: 'javascript' }
337+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
338338
})
339339

340340
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -398,7 +398,7 @@ describe('integrations', () => {
398398
metadata: MOCK_ANY,
399399
// also tests tokens not sent on llm-type spans should be 0
400400
tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 },
401-
tags: { ml_app: 'test', language: 'javascript' }
401+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
402402
})
403403

404404
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -455,7 +455,7 @@ describe('integrations', () => {
455455
inputDocuments: [{ text: 'Hello!' }],
456456
outputValue: '[1 embedding(s) returned with size 2]',
457457
metadata: MOCK_ANY,
458-
tags: { ml_app: 'test', language: 'javascript' }
458+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
459459
})
460460

461461
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -482,7 +482,7 @@ describe('integrations', () => {
482482
inputDocuments: [{ text: 'Hello!' }],
483483
outputValue: '',
484484
metadata: MOCK_ANY,
485-
tags: { ml_app: 'test', language: 'javascript' },
485+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' },
486486
error: 1,
487487
errorType: 'Error',
488488
errorMessage: MOCK_STRING,
@@ -533,7 +533,7 @@ describe('integrations', () => {
533533
inputDocuments: [{ text: 'Hello!' }, { text: 'World!' }],
534534
outputValue: '[2 embedding(s) returned with size 2]',
535535
metadata: MOCK_ANY,
536-
tags: { ml_app: 'test', language: 'javascript' }
536+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
537537
})
538538

539539
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -583,7 +583,7 @@ describe('integrations', () => {
583583
inputValue: JSON.stringify({ input: 'Can you tell me about LangSmith?' }),
584584
outputValue: 'LangSmith can help with testing in several ways.',
585585
metadata: MOCK_ANY,
586-
tags: { ml_app: 'test', language: 'javascript' }
586+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
587587
})
588588

589589
const expectedLLM = expectedLLMObsLLMSpanEvent({
@@ -601,7 +601,7 @@ describe('integrations', () => {
601601
outputMessages: [{ content: 'LangSmith can help with testing in several ways.' }],
602602
metadata: MOCK_ANY,
603603
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
604-
tags: { ml_app: 'test', language: 'javascript' }
604+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
605605
})
606606

607607
expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow)
@@ -631,7 +631,7 @@ describe('integrations', () => {
631631
inputValue: 'Hello!',
632632
outputValue: '',
633633
metadata: MOCK_ANY,
634-
tags: { ml_app: 'test', language: 'javascript' },
634+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' },
635635
error: 1,
636636
errorType: 'Error',
637637
errorMessage: MOCK_STRING,
@@ -721,7 +721,7 @@ describe('integrations', () => {
721721
name: 'langchain_core.runnables.RunnableSequence',
722722
inputValue: JSON.stringify({ person: 'Abraham Lincoln', language: 'Spanish' }),
723723
outputValue: 'Springfield, Illinois está en los Estados Unidos.',
724-
tags: { ml_app: 'test', language: 'javascript' }
724+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
725725
})
726726

727727
const expectedFirstSubWorkflow = expectedLLMObsNonLLMSpanEvent({
@@ -731,7 +731,7 @@ describe('integrations', () => {
731731
name: 'langchain_core.runnables.RunnableSequence',
732732
inputValue: JSON.stringify({ person: 'Abraham Lincoln', language: 'Spanish' }),
733733
outputValue: 'Springfield, Illinois',
734-
tags: { ml_app: 'test', language: 'javascript' }
734+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
735735
})
736736

737737
const expectedFirstLLM = expectedLLMObsLLMSpanEvent({
@@ -747,7 +747,7 @@ describe('integrations', () => {
747747
outputMessages: [{ content: 'Springfield, Illinois', role: 'assistant' }],
748748
metadata: MOCK_ANY,
749749
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
750-
tags: { ml_app: 'test', language: 'javascript' }
750+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
751751
})
752752

753753
const expectedSecondSubWorkflow = expectedLLMObsNonLLMSpanEvent({
@@ -757,7 +757,7 @@ describe('integrations', () => {
757757
name: 'langchain_core.runnables.RunnableSequence',
758758
inputValue: JSON.stringify({ language: 'Spanish', city: 'Springfield, Illinois' }),
759759
outputValue: 'Springfield, Illinois está en los Estados Unidos.',
760-
tags: { ml_app: 'test', language: 'javascript' }
760+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
761761
})
762762

763763
const expectedSecondLLM = expectedLLMObsLLMSpanEvent({
@@ -773,7 +773,7 @@ describe('integrations', () => {
773773
outputMessages: [{ content: 'Springfield, Illinois está en los Estados Unidos.', role: 'assistant' }],
774774
metadata: MOCK_ANY,
775775
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
776-
tags: { ml_app: 'test', language: 'javascript' }
776+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
777777
})
778778

779779
expect(topLevelWorkflowSpanEvent).to.deepEqualWithMockValues(expectedTopLevelWorkflow)
@@ -859,7 +859,7 @@ describe('integrations', () => {
859859
'Why did the chicken cross the road? To get to the other side!',
860860
'Why was the dog confused? It was barking up the wrong tree!'
861861
]),
862-
tags: { ml_app: 'test', language: 'javascript' }
862+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
863863
})
864864

865865
const expectedFirstLLM = expectedLLMObsLLMSpanEvent({
@@ -876,7 +876,7 @@ describe('integrations', () => {
876876
}],
877877
metadata: MOCK_ANY,
878878
tokenMetrics: { input_tokens: 37, output_tokens: 10, total_tokens: 47 },
879-
tags: { ml_app: 'test', language: 'javascript' }
879+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
880880
})
881881

882882
const expectedSecondLLM = expectedLLMObsLLMSpanEvent({
@@ -893,7 +893,7 @@ describe('integrations', () => {
893893
}],
894894
metadata: MOCK_ANY,
895895
tokenMetrics: { input_tokens: 37, output_tokens: 10, total_tokens: 47 },
896-
tags: { ml_app: 'test', language: 'javascript' }
896+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
897897
})
898898

899899
expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow)
@@ -963,7 +963,7 @@ describe('integrations', () => {
963963
content: 'Mitochondria',
964964
role: 'assistant'
965965
}),
966-
tags: { ml_app: 'test', language: 'javascript' }
966+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
967967
})
968968

969969
const expectedLLM = expectedLLMObsLLMSpanEvent({
@@ -994,7 +994,7 @@ describe('integrations', () => {
994994
outputMessages: [{ content: 'Mitochondria', role: 'assistant' }],
995995
metadata: MOCK_ANY,
996996
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
997-
tags: { ml_app: 'test', language: 'javascript' }
997+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
998998
})
999999

10001000
expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow)
@@ -1064,7 +1064,7 @@ describe('integrations', () => {
10641064
name: 'langchain_core.runnables.RunnableSequence',
10651065
inputValue: JSON.stringify({ foo: 'bar' }),
10661066
outputValue: '3 squared is 9',
1067-
tags: { ml_app: 'test', language: 'javascript' }
1067+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
10681068
})
10691069

10701070
const expectedTask = expectedLLMObsNonLLMSpanEvent({
@@ -1088,7 +1088,7 @@ describe('integrations', () => {
10881088
outputMessages: [{ content: '3 squared is 9', role: 'assistant' }],
10891089
metadata: MOCK_ANY,
10901090
tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 },
1091-
tags: { ml_app: 'test', language: 'javascript' }
1091+
tags: { ml_app: 'test', language: 'javascript', integration: 'langchain' }
10921092
})
10931093

10941094
expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow)

‎packages/dd-trace/test/llmobs/plugins/openai/openaiv3.spec.js‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ describe('integrations', () => {
9999
modelName: 'text-davinci-002',
100100
modelProvider: 'openai',
101101
metadata: {},
102-
tags: { ml_app: 'test', language: 'javascript' }
102+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' }
103103
})
104104

105105
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -156,7 +156,7 @@ describe('integrations', () => {
156156
modelName: 'gpt-3.5-turbo-0301',
157157
modelProvider: 'openai',
158158
metadata: {},
159-
tags: { ml_app: 'test', language: 'javascript' }
159+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' }
160160
})
161161

162162
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -207,7 +207,7 @@ describe('integrations', () => {
207207
modelName: 'text-embedding-ada-002-v2',
208208
modelProvider: 'openai',
209209
metadata: { encoding_format: 'float' },
210-
tags: { ml_app: 'test', language: 'javascript' }
210+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' }
211211
})
212212

213213
expect(spanEvent).to.deepEqualWithMockValues(expected)
@@ -274,7 +274,7 @@ describe('integrations', () => {
274274
]
275275
}],
276276
metadata: { function_call: 'auto' },
277-
tags: { ml_app: 'test', language: 'javascript' },
277+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' },
278278
tokenMetrics: { input_tokens: 37, output_tokens: 10, total_tokens: 47 }
279279
})
280280

@@ -311,7 +311,7 @@ describe('integrations', () => {
311311
modelName: 'gpt-3.5-turbo',
312312
modelProvider: 'openai',
313313
metadata: { max_tokens: 50 },
314-
tags: { ml_app: 'test', language: 'javascript' },
314+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' },
315315
error,
316316
errorType: error.type || error.name,
317317
errorMessage: error.message,
@@ -354,7 +354,7 @@ describe('integrations', () => {
354354
modelName: 'gpt-3.5-turbo',
355355
modelProvider: 'openai',
356356
metadata: { max_tokens: 50 },
357-
tags: { ml_app: 'test', language: 'javascript' },
357+
tags: { ml_app: 'test', language: 'javascript', integration: 'openai' },
358358
error,
359359
errorType: error.type || error.name,
360360
errorMessage: error.message,

0 commit comments

Comments
 (0)