Skip to content

Commit b874514

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
authored
feat(ai_guard): add telemetry to the SDK (#6767)
feat(ai_guard): add telemetry to the SDK
1 parent 50146d9 commit b874514

File tree

3 files changed

+81
-9
lines changed

3 files changed

+81
-9
lines changed

‎packages/dd-trace/src/aiguard/sdk.js‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,15 @@ const {
99
AI_GUARD_ACTION_TAG_KEY,
1010
AI_GUARD_BLOCKED_TAG_KEY,
1111
AI_GUARD_META_STRUCT_KEY,
12-
AI_GUARD_TOOL_NAME_TAG_KEY
12+
AI_GUARD_TOOL_NAME_TAG_KEY,
13+
AI_GUARD_TELEMETRY_REQUESTS,
14+
AI_GUARD_TELEMETRY_TRUNCATED
1315
} = require('./tags')
1416
const log = require('../log')
17+
const telemetryMetrics = require('../telemetry/metrics')
18+
const tracerVersion = require('../../../../package.json').version
19+
20+
const appsecMetrics = telemetryMetrics.manager.namespace('appsec')
1521

1622
const ALLOW = 'ALLOW'
1723

@@ -58,6 +64,9 @@ class AIGuard extends NoopAIGuard {
5864
this.#headers = {
5965
'DD-API-KEY': config.apiKey,
6066
'DD-APPLICATION-KEY': config.appKey,
67+
'DD-AI-GUARD-VERSION': tracerVersion,
68+
'DD-AI-GUARD-SOURCE': 'SDK',
69+
'DD-AI-GUARD-LANGUAGE': 'nodejs'
6170
}
6271
const endpoint = config.experimental.aiguard.endpoint || `https://app.${config.site}/api/v2/ai-guard`
6372
this.#evaluateUrl = `${endpoint}/evaluate`
@@ -70,14 +79,22 @@ class AIGuard extends NoopAIGuard {
7079

7180
#truncate (messages) {
7281
const size = Math.min(messages.length, this.#maxMessagesLength)
82+
if (messages.length > size) {
83+
appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'messages' }).inc(1)
84+
}
7385
const result = messages.slice(-size)
7486

87+
let contentTruncated = false
7588
for (let i = 0; i < size; i++) {
7689
const message = result[i]
7790
if (message.content?.length > this.#maxContentSize) {
91+
contentTruncated = true
7892
result[i] = { ...message, content: message.content.slice(0, this.#maxContentSize) }
7993
}
8094
}
95+
if (contentTruncated) {
96+
appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'content' }).inc(1)
97+
}
8198
return result
8299
}
83100

@@ -140,9 +157,11 @@ class AIGuard extends NoopAIGuard {
140157
payload,
141158
{ url: this.#evaluateUrl, headers: this.#headers, timeout: this.#timeout })
142159
} catch (e) {
143-
throw new AIGuardClientError('Unexpected error calling AI Guard service', { cause: e })
160+
appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
161+
throw new AIGuardClientError(`Unexpected error calling AI Guard service: ${e.message}`, { cause: e })
144162
}
145163
if (response.status !== 200) {
164+
appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
146165
throw new AIGuardClientError(
147166
`AI Guard service call failed, status ${response.status}`,
148167
{ errors: response.body?.errors })
@@ -157,11 +176,14 @@ class AIGuard extends NoopAIGuard {
157176
reason = attr.reason
158177
blockingEnabled = attr.is_blocking_enabled ?? false
159178
} catch (e) {
179+
appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
160180
throw new AIGuardClientError(`AI Guard service returned unexpected response : ${response.body}`, { cause: e })
161181
}
182+
const shouldBlock = block && blockingEnabled && action !== ALLOW
183+
appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { action, error: false, block: shouldBlock }).inc(1)
162184
span.setTag(AI_GUARD_ACTION_TAG_KEY, action)
163185
span.setTag(AI_GUARD_REASON_TAG_KEY, reason)
164-
if (block && blockingEnabled && action !== ALLOW) {
186+
if (shouldBlock) {
165187
span.setTag(AI_GUARD_BLOCKED_TAG_KEY, 'true')
166188
throw new AIGuardAbortError(reason)
167189
}

‎packages/dd-trace/src/aiguard/tags.js‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,8 @@ module.exports = {
77
AI_GUARD_ACTION_TAG_KEY: 'ai_guard.action',
88
AI_GUARD_REASON_TAG_KEY: 'ai_guard.reason',
99
AI_GUARD_BLOCKED_TAG_KEY: 'ai_guard.blocked',
10-
AI_GUARD_META_STRUCT_KEY: 'ai_guard'
10+
AI_GUARD_META_STRUCT_KEY: 'ai_guard',
11+
12+
AI_GUARD_TELEMETRY_REQUESTS: 'ai_guard.requests',
13+
AI_GUARD_TELEMETRY_TRUNCATED: 'ai_guard.truncated'
1114
}

‎packages/dd-trace/test/aiguard/index.spec.js‎

Lines changed: 52 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ const sinon = require('sinon')
88
const agent = require('../plugins/agent')
99
const NoopAIGuard = require('../../src/aiguard/noop')
1010
const AIGuard = require('../../src/aiguard/sdk')
11+
const tracerVersion = require('../../../../package.json').version
12+
const telemetryMetrics = require('../../src/telemetry/metrics')
13+
const appsecNamespace = telemetryMetrics.manager.namespace('appsec')
1114

1215
describe('AIGuard SDK', () => {
1316
const config = {
@@ -29,6 +32,7 @@ describe('AIGuard SDK', () => {
2932
}
3033
let tracer
3134
let aiguard
35+
let count, inc
3236

3337
const toolCall = [
3438
{ role: 'system', content: 'You are a beautiful AI assistant' },
@@ -67,21 +71,32 @@ describe('AIGuard SDK', () => {
6771
originalFetch = global.fetch
6872
global.fetch = sinon.stub()
6973

74+
inc = sinon.spy()
75+
count = sinon.stub(appsecNamespace, 'count').returns({
76+
inc
77+
})
78+
appsecNamespace.metrics.clear()
79+
7080
aiguard = new AIGuard(tracer, config)
7181

7282
return agent.load(null, [])
7383
})
7484

7585
afterEach(() => {
7686
global.fetch = originalFetch
87+
sinon.restore()
7788
agent.close()
7889
})
7990

8091
const mockFetch = (options) => {
81-
global.fetch.resolves({
82-
status: options.status ?? 200,
83-
json: sinon.stub().resolves(options.body)
84-
})
92+
if (options.error) {
93+
global.fetch.rejects(options.error)
94+
} else {
95+
global.fetch.resolves({
96+
status: options.status ?? 200,
97+
json: sinon.stub().resolves(options.body)
98+
})
99+
}
85100
}
86101

87102
const assertFetch = (messages, url) => {
@@ -96,7 +111,10 @@ describe('AIGuard SDK', () => {
96111
'Content-Type': 'application/json',
97112
'Content-Length': Buffer.byteLength(postData),
98113
'DD-API-KEY': config.apiKey,
99-
'DD-APPLICATION-KEY': config.appKey
114+
'DD-APPLICATION-KEY': config.appKey,
115+
'DD-AI-GUARD-VERSION': tracerVersion,
116+
'DD-AI-GUARD-SOURCE': 'SDK',
117+
'DD-AI-GUARD-LANGUAGE': 'nodejs'
100118
},
101119
body: postData,
102120
signal: sinon.match.instanceOf(AbortSignal)
@@ -115,6 +133,10 @@ describe('AIGuard SDK', () => {
115133
}, { rejectFirst: true })
116134
}
117135

136+
const assertTelemetry = (metric, tags) => {
137+
sinon.assert.calledWith(count, metric, tags)
138+
}
139+
118140
const testSuite = [
119141
{ action: 'ALLOW', reason: 'Go ahead' },
120142
{ action: 'DENY', reason: 'Nope' },
@@ -144,6 +166,7 @@ describe('AIGuard SDK', () => {
144166
expect(evaluation.reason).to.equal(reason)
145167
}
146168

169+
assertTelemetry('ai_guard.requests', { error: false, action, block: shouldBlock })
147170
assertFetch(messages)
148171
await assertAIGuardSpan({
149172
'ai_guard.target': target,
@@ -169,6 +192,26 @@ describe('AIGuard SDK', () => {
169192
err.name === 'AIGuardClientError' && JSON.stringify(err.errors) === JSON.stringify(errors)
170193
)
171194

195+
assertTelemetry('ai_guard.requests', { error: true })
196+
assertFetch(toolCall)
197+
await assertAIGuardSpan({
198+
'ai_guard.target': 'tool',
199+
'error.type': 'AIGuardClientError'
200+
})
201+
})
202+
203+
it('test evaluate with API exception', async () => {
204+
mockFetch({
205+
error: new Error('Boom!!!'),
206+
})
207+
208+
await rejects(
209+
() => aiguard.evaluate(toolCall),
210+
err =>
211+
err.name === 'AIGuardClientError' && err.message === 'Unexpected error calling AI Guard service: Boom!!!',
212+
)
213+
214+
assertTelemetry('ai_guard.requests', { error: true })
172215
assertFetch(toolCall)
173216
await assertAIGuardSpan({
174217
'ai_guard.target': 'tool',
@@ -184,6 +227,7 @@ describe('AIGuard SDK', () => {
184227
err => err.name === 'AIGuardClientError'
185228
)
186229

230+
assertTelemetry('ai_guard.requests', { error: true })
187231
assertFetch(toolCall)
188232
await assertAIGuardSpan({
189233
'ai_guard.target': 'tool',
@@ -199,6 +243,7 @@ describe('AIGuard SDK', () => {
199243
err => err.name === 'AIGuardClientError'
200244
)
201245

246+
assertTelemetry('ai_guard.requests', { error: true })
202247
assertFetch(toolCall)
203248
await assertAIGuardSpan({
204249
'ai_guard.target': 'tool',
@@ -225,6 +270,7 @@ describe('AIGuard SDK', () => {
225270

226271
await aiguard.evaluate(messages)
227272

273+
assertTelemetry('ai_guard.truncated', { type: 'messages' })
228274
assertFetch(messages)
229275
await assertAIGuardSpan(
230276
{ 'ai_guard.target': 'prompt', 'ai_guard.action': 'ALLOW' },
@@ -242,6 +288,7 @@ describe('AIGuard SDK', () => {
242288

243289
await aiguard.evaluate(messages)
244290

291+
assertTelemetry('ai_guard.truncated', { type: 'content' })
245292
assertFetch(messages)
246293
await assertAIGuardSpan(
247294
{ 'ai_guard.target': 'prompt', 'ai_guard.action': 'ALLOW' },

0 commit comments

Comments
 (0)