Skip to content

Commit 3f8ed41

Browse files
sabrennersabrenner
committed
fix(ai): properly handle partial experimental_telemetry configurations to wrap the correct tracer (#6319)
1 parent d71520e commit 3f8ed41

File tree

2 files changed

+155
-1
lines changed

2 files changed

+155
-1
lines changed

‎packages/datadog-instrumentations/src/ai.js

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,18 @@ function wrapWithTracer (fn) {
9898
return function () {
9999
const options = arguments[0]
100100

101-
options.experimental_telemetry ??= { isEnabled: true, tracer: noopTracer }
101+
const experimentalTelemetry = options.experimental_telemetry
102+
if (experimentalTelemetry?.isEnabled === false) {
103+
return fn.apply(this, arguments)
104+
}
105+
106+
if (experimentalTelemetry == null) {
107+
options.experimental_telemetry = { isEnabled: true, tracer: noopTracer }
108+
} else {
109+
experimentalTelemetry.isEnabled = true
110+
experimentalTelemetry.tracer ??= noopTracer
111+
}
112+
102113
wrapTracer(options.experimental_telemetry.tracer)
103114

104115
return fn.apply(this, arguments)

‎packages/datadog-plugin-ai/test/index.spec.js

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,32 @@ function getAiSdkOpenAiPackage (vercelAiVersion) {
1515
return semifies(vercelAiVersion, '>=5.0.0') ? '@ai-sdk/openai' : '@ai-sdk/[email protected]'
1616
}
1717

18+
// making a different reference from the default no-op tracer in the instrumentation
19+
// attempted to use the DD tracer provider, but it double-traces the request
20+
// in practice, there is no need to pass in the DD OTel tracer provider, so this
21+
// case shouldn't be an issue in practice
22+
const myTracer = {
23+
startActiveSpan () {
24+
const fn = arguments[arguments.length - 1]
25+
26+
const span = {
27+
spanContext () { return { traceId: '', spanId: '', traceFlags: 0 } },
28+
setAttribute () { return this },
29+
setAttributes () { return this },
30+
addEvent () { return this },
31+
addLink () { return this },
32+
addLinks () { return this },
33+
setStatus () { return this },
34+
updateName () { return this },
35+
end () { return this },
36+
isRecording () { return false },
37+
recordException () { return this }
38+
}
39+
40+
return fn(span)
41+
}
42+
}
43+
1844
describe('Plugin', () => {
1945
useEnv({
2046
OPENAI_API_KEY: '<not-a-real-key>'
@@ -38,6 +64,123 @@ describe('Plugin', () => {
3864
})
3965
})
4066

67+
describe('patching behavior with experimental_telemetry options', () => {
68+
it('should not error when `isEnabled` is false', async () => {
69+
const experimentalTelemetry = { isEnabled: false }
70+
const result = await ai.generateText({
71+
model: openai('gpt-4o-mini'),
72+
system: 'You are a helpful assistant',
73+
prompt: 'Hello, OpenAI!',
74+
maxTokens: 100,
75+
temperature: 0.5,
76+
experimental_telemetry: experimentalTelemetry
77+
})
78+
79+
assert.ok(result.text, 'Expected result to be truthy')
80+
assert.ok(experimentalTelemetry.tracer == null, 'Tracer should not be set by default')
81+
})
82+
83+
it('should not error when a `tracer` is not passed in', async () => {
84+
const checkTraces = agent.assertSomeTraces(traces => {
85+
const generateTextSpan = traces[0][0]
86+
const doGenerateSpan = traces[0][1]
87+
88+
assert.strictEqual(generateTextSpan.name, 'ai.generateText')
89+
assert.strictEqual(generateTextSpan.resource, 'ai.generateText')
90+
assert.strictEqual(generateTextSpan.meta['ai.request.model'], 'gpt-4o-mini')
91+
assert.strictEqual(generateTextSpan.meta['ai.request.model_provider'], 'openai')
92+
93+
assert.strictEqual(doGenerateSpan.name, 'ai.generateText.doGenerate')
94+
assert.strictEqual(doGenerateSpan.resource, 'ai.generateText.doGenerate')
95+
assert.strictEqual(doGenerateSpan.meta['ai.request.model'], 'gpt-4o-mini')
96+
assert.strictEqual(doGenerateSpan.meta['ai.request.model_provider'], 'openai')
97+
})
98+
99+
const experimentalTelemetry = { isEnabled: true }
100+
101+
const result = await ai.generateText({
102+
model: openai('gpt-4o-mini'),
103+
system: 'You are a helpful assistant',
104+
prompt: 'Hello, OpenAI!',
105+
maxTokens: 100,
106+
temperature: 0.5,
107+
experimental_telemetry: experimentalTelemetry
108+
})
109+
110+
assert.ok(result.text, 'Expected result to be truthy')
111+
assert.ok(experimentalTelemetry.tracer != null, 'Tracer should be set when `isEnabled` is true')
112+
113+
await checkTraces
114+
})
115+
116+
it('should not error when only a `tracer` is not passed in', async () => {
117+
const checkTraces = agent.assertSomeTraces(traces => {
118+
const generateTextSpan = traces[0][0]
119+
const doGenerateSpan = traces[0][1]
120+
121+
assert.strictEqual(generateTextSpan.name, 'ai.generateText')
122+
assert.strictEqual(generateTextSpan.resource, 'ai.generateText')
123+
assert.strictEqual(generateTextSpan.meta['ai.request.model'], 'gpt-4o-mini')
124+
assert.strictEqual(generateTextSpan.meta['ai.request.model_provider'], 'openai')
125+
126+
assert.strictEqual(doGenerateSpan.name, 'ai.generateText.doGenerate')
127+
assert.strictEqual(doGenerateSpan.resource, 'ai.generateText.doGenerate')
128+
assert.strictEqual(doGenerateSpan.meta['ai.request.model'], 'gpt-4o-mini')
129+
assert.strictEqual(doGenerateSpan.meta['ai.request.model_provider'], 'openai')
130+
})
131+
132+
const experimentalTelemetry = { tracer: myTracer }
133+
134+
const result = await ai.generateText({
135+
model: openai('gpt-4o-mini'),
136+
system: 'You are a helpful assistant',
137+
prompt: 'Hello, OpenAI!',
138+
maxTokens: 100,
139+
temperature: 0.5,
140+
experimental_telemetry: experimentalTelemetry
141+
})
142+
143+
assert.ok(result.text, 'Expected result to be truthy')
144+
assert.ok(experimentalTelemetry.isEnabled, 'isEnabled should be set to true')
145+
assert.ok(experimentalTelemetry.tracer === myTracer, 'Tracer should be set when `isEnabled` is true')
146+
147+
await checkTraces
148+
})
149+
150+
it('should use the passed in `tracer`', async () => {
151+
const checkTraces = agent.assertSomeTraces(traces => {
152+
const generateTextSpan = traces[0][0]
153+
const doGenerateSpan = traces[0][1]
154+
155+
assert.strictEqual(generateTextSpan.name, 'ai.generateText')
156+
assert.strictEqual(generateTextSpan.resource, 'ai.generateText')
157+
assert.strictEqual(generateTextSpan.meta['ai.request.model'], 'gpt-4o-mini')
158+
assert.strictEqual(generateTextSpan.meta['ai.request.model_provider'], 'openai')
159+
160+
assert.strictEqual(doGenerateSpan.name, 'ai.generateText.doGenerate')
161+
assert.strictEqual(doGenerateSpan.resource, 'ai.generateText.doGenerate')
162+
assert.strictEqual(doGenerateSpan.meta['ai.request.model'], 'gpt-4o-mini')
163+
assert.strictEqual(doGenerateSpan.meta['ai.request.model_provider'], 'openai')
164+
})
165+
166+
const experimentalTelemetry = { isEnabled: true, tracer: myTracer }
167+
168+
const result = await ai.generateText({
169+
model: openai('gpt-4o-mini'),
170+
system: 'You are a helpful assistant',
171+
prompt: 'Hello, OpenAI!',
172+
maxTokens: 100,
173+
temperature: 0.5,
174+
experimental_telemetry: experimentalTelemetry
175+
})
176+
177+
assert.ok(result.text, 'Expected result to be truthy')
178+
assert.ok(experimentalTelemetry.tracer === myTracer, 'Tracer should not override provided tracer')
179+
180+
await checkTraces
181+
})
182+
})
183+
41184
it('creates a span for generateText', async () => {
42185
const checkTraces = agent.assertSomeTraces(traces => {
43186
const generateTextSpan = traces[0][0]

0 commit comments

Comments
 (0)