Skip to content

Commit 6244195

Browse files
sabrennersabrenner
authored
feat(ai, llmobs): properly support ToolLoopAgent via existing patching (#7571)
wip - mostly working everything works but later 5.x versions in tests test fixes add tests for toolloopagent generating spans remove load publish remove hook file instead of deleting, with comment address review comment use sets instead of arrays for orchestrion supported version checks Merge branch 'master' into sabrenner/vercel-ai-use-orchestrion Co-authored-by: sam.brenner <[email protected]>
1 parent 05f01a0 commit 6244195

File tree

12 files changed

+959
-105
lines changed

12 files changed

+959
-105
lines changed

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

Lines changed: 54 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -2,45 +2,11 @@
22

33
const { channel, tracingChannel } = require('dc-polyfill')
44
const shimmer = require('../../datadog-shimmer')
5-
const { addHook } = require('./helpers/instrument')
6-
7-
const toolCreationChannel = channel('dd-trace:vercel-ai:tool')
8-
9-
const TRACED_FUNCTIONS = {
10-
generateText: wrapWithTracer,
11-
streamText: wrapWithTracer,
12-
generateObject: wrapWithTracer,
13-
streamObject: wrapWithTracer,
14-
embed: wrapWithTracer,
15-
embedMany: wrapWithTracer,
16-
tool: wrapTool,
17-
}
5+
const { addHook, getHooks } = require('./helpers/instrument')
186

197
const vercelAiTracingChannel = tracingChannel('dd-trace:vercel-ai')
208
const vercelAiSpanSetAttributesChannel = channel('dd-trace:vercel-ai:span:setAttributes')
219

22-
const noopTracer = {
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-
4410
const tracers = new WeakSet()
4511

4612
function wrapTracer (tracer) {
@@ -63,29 +29,35 @@ function wrapTracer (tracer) {
6329

6430
arguments[arguments.length - 1] = shimmer.wrapFunction(cb, function (originalCb) {
6531
return function (span) {
66-
shimmer.wrap(span, 'end', function (spanEnd) {
32+
// the below is necessary in the case that the span is vercel ai's noopSpan.
33+
// while we don't want to patch the noopSpan more than once, we do want to treat each as a
34+
// fresh instance. However, this is really not necessary for non-noop spans, but not sure
35+
// how to differentiate.
36+
const freshSpan = Object.create(span) // TODO: does this cause memory leaks?
37+
38+
shimmer.wrap(freshSpan, 'end', function (spanEnd) {
6739
return function () {
6840
vercelAiTracingChannel.asyncEnd.publish(ctx)
6941
return spanEnd.apply(this, arguments)
7042
}
7143
})
7244

73-
shimmer.wrap(span, 'setAttributes', function (setAttributes) {
45+
shimmer.wrap(freshSpan, 'setAttributes', function (setAttributes) {
7446
return function (attributes) {
7547
vercelAiSpanSetAttributesChannel.publish({ ctx, attributes })
7648
return setAttributes.apply(this, arguments)
7749
}
7850
})
7951

80-
shimmer.wrap(span, 'recordException', function (recordException) {
52+
shimmer.wrap(freshSpan, 'recordException', function (recordException) {
8153
return function (exception) {
8254
ctx.error = exception
8355
vercelAiTracingChannel.error.publish(ctx)
8456
return recordException.apply(this, arguments)
8557
}
8658
})
8759

88-
return originalCb.apply(this, arguments)
60+
return originalCb.call(this, freshSpan)
8961
}
9062
})
9163

@@ -98,58 +70,50 @@ function wrapTracer (tracer) {
9870
})
9971
}
10072

101-
function wrapWithTracer (fn) {
102-
return function () {
103-
const options = arguments[0]
104-
105-
const experimentalTelemetry = options.experimental_telemetry
106-
if (experimentalTelemetry?.isEnabled === false) {
107-
return fn.apply(this, arguments)
108-
}
109-
110-
if (experimentalTelemetry == null) {
111-
options.experimental_telemetry = { isEnabled: true, tracer: noopTracer }
112-
} else {
113-
experimentalTelemetry.isEnabled = true
114-
experimentalTelemetry.tracer ??= noopTracer
115-
}
116-
117-
wrapTracer(options.experimental_telemetry.tracer)
118-
119-
return fn.apply(this, arguments)
73+
for (const hook of getHooks('ai')) {
74+
if (hook.file === 'dist/index.js') {
75+
// if not removed, the below hook will never match correctly
76+
// however, it is still needed in the orchestrion definition
77+
hook.file = null
12078
}
121-
}
122-
123-
function wrapTool (tool) {
124-
return function () {
125-
const args = arguments[0]
126-
toolCreationChannel.publish(args)
127-
128-
return tool.apply(this, arguments)
129-
}
130-
}
13179

132-
// CJS exports
133-
addHook({
134-
name: 'ai',
135-
versions: ['>=4.0.0'],
136-
}, exports => {
137-
for (const [fnName, patchingFn] of Object.entries(TRACED_FUNCTIONS)) {
138-
exports = shimmer.wrap(exports, fnName, patchingFn, { replaceGetter: true })
139-
}
80+
addHook(hook, exports => {
81+
const getTracerChannel = tracingChannel('orchestrion:ai:getTracer')
82+
getTracerChannel.subscribe({
83+
end (ctx) {
84+
const { arguments: args, result: tracer } = ctx
85+
const { isEnabled } = args[0] ?? {}
14086

141-
return exports
142-
})
143-
144-
// ESM exports
145-
addHook({
146-
name: 'ai',
147-
versions: ['>=4.0.0'],
148-
file: 'dist/index.mjs',
149-
}, exports => {
150-
for (const [fnName, patchingFn] of Object.entries(TRACED_FUNCTIONS)) {
151-
exports = shimmer.wrap(exports, fnName, patchingFn, { replaceGetter: true })
152-
}
87+
if (isEnabled !== false) {
88+
wrapTracer(tracer)
89+
}
90+
},
91+
})
92+
93+
/**
94+
* We patch this function to ensure that the telemetry attributes/tags are set always,
95+
* even when telemetry options are not specified. This is to ensure easy use of this integration.
96+
*
97+
* If it is explicitly disabled, however, we will not change the options.
98+
*/
99+
const selectTelemetryAttributesChannel = tracingChannel('orchestrion:ai:selectTelemetryAttributes')
100+
selectTelemetryAttributesChannel.subscribe({
101+
start (ctx) {
102+
const { arguments: args } = ctx
103+
const options = args[0]
104+
105+
if (options.telemetry?.isEnabled !== false) {
106+
args[0] = {
107+
...options,
108+
telemetry: {
109+
...options.telemetry,
110+
isEnabled: true,
111+
},
112+
}
113+
}
114+
},
115+
})
153116

154-
return exports
155-
})
117+
return exports
118+
})
119+
}

‎packages/datadog-instrumentations/src/helpers/rewriter/index.js‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ const instrumentations = require('./instrumentations')
4343

4444
const NODE_OPTIONS = getEnvironmentVariable('NODE_OPTIONS')
4545

46+
/** @type {Record<string, Set<string>>} map of module base name to supported function query versions */
4647
const supported = {}
4748
const disabled = new Set()
4849

@@ -101,19 +102,21 @@ function disable (instrumentation) {
101102
function satisfies (filename, filePath, versions) {
102103
const [basename] = filename.split(filePath)
103104

104-
if (supported[basename] === undefined) {
105+
supported[basename] ??= new Set()
106+
107+
if (!supported[basename].has(versions)) {
105108
try {
106109
const pkg = JSON.parse(readFileSync(
107110
join(basename, 'package.json'), 'utf8'
108111
))
109112

110-
supported[basename] = semifies(pkg.version, versions)
111-
} catch {
112-
supported[basename] = false
113-
}
113+
if (semifies(pkg.version, versions)) {
114+
supported[basename].add(versions)
115+
}
116+
} catch {}
114117
}
115118

116-
return supported[basename]
119+
return supported[basename].has(versions)
117120
}
118121

119122
// TODO: Support index

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

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
'use strict'
2+
3+
module.exports = [
4+
// getTracer - for patching tracer
5+
{
6+
module: {
7+
name: 'ai',
8+
versionRange: '>=4.0.0',
9+
filePath: 'dist/index.js',
10+
},
11+
functionQuery: {
12+
functionName: 'getTracer',
13+
kind: 'Sync',
14+
},
15+
channelName: 'getTracer',
16+
},
17+
{
18+
module: {
19+
name: 'ai',
20+
versionRange: '>=4.0.0',
21+
filePath: 'dist/index.mjs',
22+
},
23+
functionQuery: {
24+
functionName: 'getTracer',
25+
kind: 'Sync',
26+
},
27+
channelName: 'getTracer',
28+
},
29+
// selectTelemetryAttributes - makes sure we set isEnabled properly
30+
{
31+
module: {
32+
name: 'ai',
33+
versionRange: '>=4.0.0 <6.0.0',
34+
filePath: 'dist/index.js',
35+
},
36+
functionQuery: {
37+
functionName: 'selectTelemetryAttributes',
38+
kind: 'Sync',
39+
},
40+
channelName: 'selectTelemetryAttributes',
41+
},
42+
{
43+
module: {
44+
name: 'ai',
45+
versionRange: '>=4.0.0 <6.0.0',
46+
filePath: 'dist/index.mjs',
47+
},
48+
functionQuery: {
49+
functionName: 'selectTelemetryAttributes',
50+
kind: 'Sync',
51+
},
52+
channelName: 'selectTelemetryAttributes',
53+
},
54+
{
55+
module: {
56+
name: 'ai',
57+
versionRange: '>=6.0.0',
58+
filePath: 'dist/index.js',
59+
},
60+
functionQuery: {
61+
functionName: 'selectTelemetryAttributes',
62+
kind: 'Async',
63+
},
64+
channelName: 'selectTelemetryAttributes',
65+
},
66+
{
67+
module: {
68+
name: 'ai',
69+
versionRange: '>=6.0.0',
70+
filePath: 'dist/index.mjs',
71+
},
72+
functionQuery: {
73+
functionName: 'selectTelemetryAttributes',
74+
kind: 'Async',
75+
},
76+
channelName: 'selectTelemetryAttributes',
77+
},
78+
// tool
79+
{
80+
module: {
81+
name: 'ai',
82+
versionRange: '>=4.0.0',
83+
filePath: 'dist/index.js',
84+
},
85+
functionQuery: {
86+
functionName: 'tool',
87+
kind: 'Sync',
88+
},
89+
channelName: 'tool',
90+
},
91+
{
92+
module: {
93+
name: 'ai',
94+
versionRange: '>=4.0.0',
95+
filePath: 'dist/index.mjs',
96+
},
97+
functionQuery: {
98+
functionName: 'tool',
99+
kind: 'Sync',
100+
},
101+
channelName: 'tool',
102+
},
103+
]

‎packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/index.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
'use strict'
22

33
module.exports = [
4-
...require('./langchain'),
4+
...require('./ai'),
55
...require('./bullmq'),
6+
...require('./langchain'),
67
]

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,13 @@ const { NODE_MAJOR } = require('../../../version')
1212
const range = NODE_MAJOR < 22 ? '>=4.0.2' : '>=4.0.0'
1313

1414
function getAiSdkOpenAiPackage (vercelAiVersion) {
15-
return semifies(vercelAiVersion, '>=5.0.0') ? '@ai-sdk/openai' : '@ai-sdk/[email protected]'
15+
if (semifies(vercelAiVersion, '>=6.0.0')) {
16+
return '@ai-sdk/openai'
17+
} else if (semifies(vercelAiVersion, '>=5.0.0')) {
18+
return '@ai-sdk/[email protected]'
19+
} else {
20+
return '@ai-sdk/[email protected]'
21+
}
1622
}
1723

1824
// making a different reference from the default no-op tracer in the instrumentation
@@ -116,7 +122,6 @@ describe('Plugin', () => {
116122
})
117123

118124
assert.ok(result.text, 'Expected result to be truthy')
119-
assert.ok(experimentalTelemetry.tracer != null, 'Tracer should be set when `isEnabled` is true')
120125

121126
await checkTraces
122127
})
@@ -157,7 +162,6 @@ describe('Plugin', () => {
157162
})
158163

159164
assert.ok(result.text, 'Expected result to be truthy')
160-
assert.ok(experimentalTelemetry.isEnabled, 'isEnabled should be set to true')
161165
assert.ok(experimentalTelemetry.tracer === myTracer, 'Tracer should be set when `isEnabled` is true')
162166

163167
await checkTraces

‎packages/dd-trace/src/llmobs/plugins/ai/index.js‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ const { channel } = require('dc-polyfill')
44
const BaseLLMObsPlugin = require('../base')
55
const { getModelProvider } = require('../../../../../datadog-plugin-ai/src/utils')
66

7-
const toolCreationCh = channel('dd-trace:vercel-ai:tool')
7+
const toolCreationCh = channel('tracing:orchestrion:ai:tool:start')
88
const setAttributesCh = channel('dd-trace:vercel-ai:span:setAttributes')
99

1010
const { MODEL_NAME, MODEL_PROVIDER, NAME } = require('../../constants/tags')
@@ -94,8 +94,10 @@ class VercelAILLMObsPlugin extends BaseLLMObsPlugin {
9494

9595
this.#toolCallIdsToName = {}
9696
this.#availableTools = new Set()
97-
toolCreationCh.subscribe(toolArgs => {
98-
this.#availableTools.add(toolArgs)
97+
toolCreationCh.subscribe(ctx => {
98+
const toolArgs = ctx.arguments
99+
const tool = toolArgs[0] ?? {}
100+
this.#availableTools.add(tool)
99101
})
100102

101103
setAttributesCh.subscribe(({ ctx, attributes }) => {

0 commit comments

Comments
 (0)