feat(ai_guard): add telemetry to the SDK

Author: manuel-alvarez-alvarez

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

@@ -9,9 +9,14 @@ const {
   AI_GUARD_ACTION_TAG_KEY,
   AI_GUARD_BLOCKED_TAG_KEY,
   AI_GUARD_META_STRUCT_KEY,
-  AI_GUARD_TOOL_NAME_TAG_KEY
+  AI_GUARD_TOOL_NAME_TAG_KEY,
+  AI_GUARD_TELEMETRY_REQUESTS,
+  AI_GUARD_TELEMETRY_TRUNCATED
 } = require('./tags')
 const log = require('../log')
+const telemetryMetrics = require('../telemetry/metrics')
+
+const appsecMetrics = telemetryMetrics.manager.namespace('appsec')
 
 const ALLOW = 'ALLOW'
 
@@ -70,14 +75,22 @@ class AIGuard extends NoopAIGuard {
 
   #truncate (messages) {
     const size = Math.min(messages.length, this.#maxMessagesLength)
+    if (messages.length > size) {
+      appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'messages' }).inc(1)
+    }
     const result = messages.slice(-size)
 
+    let contentTruncated = false
     for (let i = 0; i < size; i++) {
       const message = result[i]
       if (message.content?.length > this.#maxContentSize) {
+        contentTruncated = true
         result[i] = { ...message, content: message.content.slice(0, this.#maxContentSize) }
       }
     }
+    if (contentTruncated) {
+      appsecMetrics.count(AI_GUARD_TELEMETRY_TRUNCATED, { type: 'content' }).inc(1)
+    }
     return result
   }
 
@@ -140,9 +153,11 @@ class AIGuard extends NoopAIGuard {
           payload,
           { url: this.#evaluateUrl, headers: this.#headers, timeout: this.#timeout })
       } catch (e) {
-        throw new AIGuardClientError('Unexpected error calling AI Guard service', { cause: e })
+        appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
+        throw new AIGuardClientError(`Unexpected error calling AI Guard service: ${e.message}`, { cause: e })
       }
       if (response.status !== 200) {
+        appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
         throw new AIGuardClientError(
           `AI Guard service call failed, status ${response.status}`,
           { errors: response.body?.errors })
@@ -157,8 +172,10 @@ class AIGuard extends NoopAIGuard {
         reason = attr.reason
         blockingEnabled = attr.is_blocking_enabled ?? false
       } catch (e) {
+        appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
         throw new AIGuardClientError(`AI Guard service returned unexpected response : ${response.body}`, { cause: e })
       }
+      appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { action, error: false, block: block && blockingEnabled }).inc(1)
       span.setTag(AI_GUARD_ACTION_TAG_KEY, action)
       span.setTag(AI_GUARD_REASON_TAG_KEY, reason)
       if (block && blockingEnabled && action !== ALLOW) {

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

@@ -7,5 +7,8 @@ module.exports = {
   AI_GUARD_ACTION_TAG_KEY: 'ai_guard.action',
   AI_GUARD_REASON_TAG_KEY: 'ai_guard.reason',
   AI_GUARD_BLOCKED_TAG_KEY: 'ai_guard.blocked',
-  AI_GUARD_META_STRUCT_KEY: 'ai_guard'
+  AI_GUARD_META_STRUCT_KEY: 'ai_guard',
+
+  AI_GUARD_TELEMETRY_REQUESTS: 'ai_guard.requests',
+  AI_GUARD_TELEMETRY_TRUNCATED: 'ai_guard.truncated'
 }

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

@@ -8,6 +8,8 @@ const sinon = require('sinon')
 const agent = require('../plugins/agent')
 const NoopAIGuard = require('../../src/aiguard/noop')
 const AIGuard = require('../../src/aiguard/sdk')
+const telemetryMetrics = require('../../src/telemetry/metrics')
+const appsecNamespace = telemetryMetrics.manager.namespace('appsec')
 
 describe('AIGuard SDK', () => {
   const config = {
@@ -29,6 +31,7 @@ describe('AIGuard SDK', () => {
   }
   let tracer
   let aiguard
+  let count, inc
 
   const toolCall = [
     { role: 'system', content: 'You are a beautiful AI assistant' },
@@ -67,13 +70,20 @@ describe('AIGuard SDK', () => {
     originalFetch = global.fetch
     global.fetch = sinon.stub()
 
+    inc = sinon.spy()
+    count = sinon.stub(appsecNamespace, 'count').returns({
+      inc
+    })
+    appsecNamespace.metrics.clear()
+
     aiguard = new AIGuard(tracer, config)
 
     return agent.load(null, [])
   })
 
   afterEach(() => {
     global.fetch = originalFetch
+    sinon.restore()
     agent.close()
   })
 
@@ -115,6 +125,10 @@ describe('AIGuard SDK', () => {
     }, { rejectFirst: true })
   }
 
+  const assertTelemetry = (metric, tags) => {
+    sinon.assert.calledWith(count, metric, tags)
+  }
+
   const testSuite = [
     { action: 'ALLOW', reason: 'Go ahead' },
     { action: 'DENY', reason: 'Nope' },
@@ -144,6 +158,7 @@ describe('AIGuard SDK', () => {
         expect(evaluation.reason).to.equal(reason)
       }
 
+      assertTelemetry('ai_guard.requests', { error: false, action, block: blocking })
       assertFetch(messages)
       await assertAIGuardSpan({
         'ai_guard.target': target,
@@ -169,6 +184,7 @@ describe('AIGuard SDK', () => {
         err.name === 'AIGuardClientError' && JSON.stringify(err.errors) === JSON.stringify(errors)
     )
 
+    assertTelemetry('ai_guard.requests', { error: true })
     assertFetch(toolCall)
     await assertAIGuardSpan({
       'ai_guard.target': 'tool',
@@ -184,6 +200,7 @@ describe('AIGuard SDK', () => {
       err => err.name === 'AIGuardClientError'
     )
 
+    assertTelemetry('ai_guard.requests', { error: true })
     assertFetch(toolCall)
     await assertAIGuardSpan({
       'ai_guard.target': 'tool',
@@ -199,6 +216,7 @@ describe('AIGuard SDK', () => {
       err => err.name === 'AIGuardClientError'
     )
 
+    assertTelemetry('ai_guard.requests', { error: true })
     assertFetch(toolCall)
     await assertAIGuardSpan({
       'ai_guard.target': 'tool',
@@ -225,6 +243,7 @@ describe('AIGuard SDK', () => {
 
     await aiguard.evaluate(messages)
 
+    assertTelemetry('ai_guard.truncated', { type: 'messages' })
     assertFetch(messages)
     await assertAIGuardSpan(
       { 'ai_guard.target': 'prompt', 'ai_guard.action': 'ALLOW' },
@@ -242,6 +261,7 @@ describe('AIGuard SDK', () => {
 
     await aiguard.evaluate(messages)
 
+    assertTelemetry('ai_guard.truncated', { type: 'content' })
     assertFetch(messages)
     await assertAIGuardSpan(
       { 'ai_guard.target': 'prompt', 'ai_guard.action': 'ALLOW' },