Expose tag probabilities in SDK responses

docs/test.ts

@@ -745,11 +745,11 @@ aiguard.evaluate([
     ],
   }
 ]).then(result => {
-  result.action && result.reason && result.tags
+  result.action && result.reason && result.tags && result.tagProbabilities && result.sds
 })
 
 aiguard.evaluate([
   { role: 'tool', tool_call_id: 'call_1', content: '5' },
 ]).then(result => {
-  result.action && result.reason && result.tags
+  result.action && result.reason && result.tags && result.tagProbabilities && result.sds
 })

index.d.ts

@@ -1847,6 +1847,10 @@ declare namespace tracer {
        * List of tags associated with the evaluation (e.g. indirect-prompt-injection)
        */
       tags: string[];
+      /**
+       * Dictionary of tag probabilities (e.g. { indirect-prompt-injection: 0.2, jailbreak-attempt: 0.8 })
+       */
+      tagProbabilities: { [key: string]: number }
       /**
        * Sensitive Data Scanner findings from the evaluation.
        */
@@ -1866,6 +1870,10 @@ declare namespace tracer {
        * List of tags associated with the evaluation (e.g. indirect-prompt-injection)
        */
       tags: string[];
+      /**
+       * Dictionary of tag probabilities (e.g. [indirect-prompt-injection: 0.2, jailbreak-attempt: 0.8])
+       */
+      tagProbabilities: { [key: string]: number }
       /**
        * Sensitive Data Scanner findings from the evaluation.
        */

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

@@ -25,11 +25,12 @@ const appsecMetrics = telemetryMetrics.manager.namespace('appsec')
 const ALLOW = 'ALLOW'
 
 class AIGuardAbortError extends Error {
-  constructor (reason, tags, sds) {
+  constructor (reason, tags, tagProbs, sds) {
     super(reason)
     this.name = 'AIGuardAbortError'
     this.reason = reason
     this.tags = tags
+    this.tagProbabilities = tagProbs
     this.sds = sds || []
   }
 }
@@ -188,7 +189,7 @@ class AIGuard extends NoopAIGuard {
           `AI Guard service call failed, status ${response.status}`,
           { errors: response.body?.errors })
       }
-      let action, reason, tags, sdsFindings, blockingEnabled
+      let action, reason, tags, sdsFindings, tagProbabilities, blockingEnabled
       try {
         const attr = response.body.data.attributes
         if (!attr.action) {
@@ -198,6 +199,7 @@ class AIGuard extends NoopAIGuard {
         reason = attr.reason
         tags = attr.tags
         sdsFindings = attr.sds_findings || []
+        tagProbabilities = attr.tagProbs || {}
         blockingEnabled = attr.is_blocking_enabled ?? false
       } catch (e) {
         appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
@@ -215,11 +217,14 @@ class AIGuard extends NoopAIGuard {
       if (sdsFindings?.length > 0) {
         metaStruct.sds = sdsFindings
       }
+      if (Object.keys(tagProbabilities).length > 0) {
+        metaStruct.tag_probs = tagProbabilities
+      }
       if (shouldBlock) {
         span.setTag(AI_GUARD_BLOCKED_TAG_KEY, 'true')
-        throw new AIGuardAbortError(reason, tags, sdsFindings)
+        throw new AIGuardAbortError(reason, tags, tagProbabilities, sdsFindings)
       }
-      return { action, reason, tags, sds: sdsFindings }
+      return { action, reason, tags, tagProbabilities, sds: sdsFindings }
     })
   }
 }

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

@@ -144,9 +144,9 @@ describe('AIGuard SDK', () => {
   }
 
   const testSuite = [
-    { action: 'ALLOW', reason: 'Go ahead', tags: [] },
-    { action: 'DENY', reason: 'Nope', tags: ['deny_everything', 'test_deny'] },
-    { action: 'ABORT', reason: 'Kill it with fire', tags: ['alarm_tag', 'abort_everything'] },
+    { action: 'ALLOW', reason: 'Go ahead', tagProbs: {} },
+    { action: 'DENY', reason: 'Nope', tagProbs: { deny_everything: 0.8, test_deny: 0.2 } },
+    { action: 'ABORT', reason: 'Kill it with fire', tagProbs: { alarm_tag: 0.3, abort_everything: 0.7 } },
   ].flatMap(r => [
     { ...r, blocking: true },
     { ...r, blocking: false },
@@ -156,23 +156,25 @@ describe('AIGuard SDK', () => {
     { ...r, suite: 'prompt', target: 'prompt', messages: prompt },
   ])
 
-  for (const { action, reason, tags, blocking, suite, target, messages } of testSuite) {
+  for (const { action, reason, tagProbs, blocking, suite, target, messages } of testSuite) {
+    const tags = Object.keys(tagProbs)
     it(`test evaluate '${suite}' with ${action} action (blocking: ${blocking})`, async () => {
-      mockFetch({ body: { data: { attributes: { action, reason, tags, is_blocking_enabled: blocking } } } })
+      mockFetch({ body: { data: { attributes: { action, reason, tags, tagProbs, is_blocking_enabled: blocking } } } })
       const shouldBlock = action !== 'ALLOW' && blocking
 
       if (shouldBlock) {
         await rejects(
           () => aiguard.evaluate(messages, { block: true }),
           err => err.name === 'AIGuardAbortError' && err.reason === reason && err.tags === tags &&
-            JSON.stringify(err.sds) === '[]'
+             err.tagProbabilities === tagProbs && JSON.stringify(err.sds) === '[]'
         )
       } else {
         const evaluation = await aiguard.evaluate(messages, { block: true })
         assert.strictEqual(evaluation.action, action)
         assert.strictEqual(evaluation.reason, reason)
-        if (tags) {
+        if (tagProbs) {
           assert.strictEqual(evaluation.tags, tags)
+          assert.strictEqual(evaluation.tagProbabilities, tagProbs)
         }
         assert.deepStrictEqual(evaluation.sds, [])
       }
@@ -189,6 +191,7 @@ describe('AIGuard SDK', () => {
       {
         messages,
         ...(tags.length > 0 ? { attack_categories: tags } : {}),
+        ...(Object.keys(tagProbs).length > 0 ? { tag_probs: tagProbs } : {}),
       })
     })
   }