aws_bedrock: add mapping for gen_ai.guardrail_id (elastic#11858)

This allows detection of a guardrail having been run by using the query gen_ai.guardrail_id:*.
harnish-elastic · Nov 25, 2024 · ef28635 · ef28635
1 parent 29e226c
commit ef28635
Show file tree

Hide file tree

Showing 7 changed files with 696 additions and 4 deletions.
diff --git a/packages/aws_bedrock/changelog.yml b/packages/aws_bedrock/changelog.yml
@@ -1,3 +1,8 @@
+- version: "0.13.0"
+  changes:
+    - description: Add mapping for `gen_ai.guardrail_id` from `output.outputBodyJson.trace.guardrail.inputAssessment`.
+      type: enhancement
+      link: https://github.com/elastic/integrations/pull/11858
 - version: "0.12.0"
   changes:
     - description: Support configuring the Owning Account

diff --git a/packages/aws_bedrock/data_stream/invocation/_dev/test/pipeline/test-aws-bedrock.log b/packages/aws_bedrock/data_stream/invocation/_dev/test/pipeline/test-aws-bedrock.log
diff --git a/.../aws_bedrock/data_stream/invocation/_dev/test/pipeline/test-aws-bedrock.log-expected.json b/.../aws_bedrock/data_stream/invocation/_dev/test/pipeline/test-aws-bedrock.log-expected.json
diff --git a/packages/aws_bedrock/data_stream/invocation/elasticsearch/ingest_pipeline/default.yml b/packages/aws_bedrock/data_stream/invocation/elasticsearch/ingest_pipeline/default.yml
@@ -355,8 +355,10 @@ processors:
         def confidences = new HashSet();
         def violations = new HashSet();
         def matches = new HashSet();
+        def guardrailIDs = new HashSet();
+        ctx._debug = new ArrayList();
         for (def e: body) {
-          if (e.amazon_bedrock_guardrail_action == 'INTERVENED' || e.stop_reason == 'guardrail_intervened') {
+          if (e.amazon_bedrock_guardrail_action == 'INTERVENED' || e.stop_reason == 'guardrail_intervened' || e.stop_reason == 'end_turn') {
             if (ctx.gen_ai == null) {
               ctx.gen_ai = new HashMap();
             }
@@ -366,10 +368,16 @@ processors:
             if (ctx.gen_ai.compliance == null) {
               ctx.gen_ai.compliance = new HashMap();
             }
-            ctx.gen_ai.compliance.violation_detected = true;
-            ctx.event.outcome = 'failure';
+            if (e.stop_reason != 'end_turn') {
+              ctx.gen_ai.compliance.violation_detected = true;
+              ctx.event.outcome = 'failure';
+            }
             if (e.trace?.guardrail?.inputAssessment instanceof HashMap) {
               for (def o: e.trace.guardrail.inputAssessment.entrySet()) {
+                guardrailIDs.add(o.getKey());
+                if (e.stop_reason == 'end_turn') {
+                  continue;
+                }
                 for (def i: o.getValue().entrySet()) {
                   String policy_name = i.getKey();
                   if (!policy_name.endsWith('_policy')) {
@@ -394,6 +402,10 @@ processors:
               }
             } else if (e.amazon_bedrock_trace?.guardrail?.input instanceof HashMap) {
               for (def o: e.amazon_bedrock_trace.guardrail.input.entrySet()) {
+                guardrailIDs.add(o.getKey());
+                if (e.stop_reason == 'end_turn') {
+                  continue;
+                }
                 for (def i: o.getValue().entrySet()) {
                   String policy_name = i.getKey();
                   if (!policy_name.endsWith('_policy')) {
@@ -419,6 +431,13 @@ processors:
             }
           }
         }
+        if (guardrailIDs.size() != 0) {
+          ctx.gen_ai.guardrail_id = new ArrayList();
+          for (def e: guardrailIDs) {
+            ctx.gen_ai.guardrail_id.add(e);
+          }
+          Collections.sort(ctx.gen_ai.guardrail_id);
+        }
         if (actions.size() != 0) {
           ctx.gen_ai.policy.action = new ArrayList();
           for (def e: actions) {

diff --git a/packages/aws_bedrock/data_stream/invocation/fields/fields.yml b/packages/aws_bedrock/data_stream/invocation/fields/fields.yml
@@ -61,6 +61,9 @@
       type: keyword
     - name: task_type
       type: keyword
+- name: gen_ai.guardrail_id
+  type: keyword
+  description: Guardrail ID if a guardrail was executed.
 - name: gen_ai.prompt
   type: text
   description: The full text of the user's request to the gen_ai.

diff --git a/packages/aws_bedrock/docs/README.md b/packages/aws_bedrock/docs/README.md
@@ -176,6 +176,7 @@ list log events from the specified log group.
 | gen_ai.compliance.response_triggered | Lists compliance-related filters that were triggered during the processing of the response, such as data privacy filters or regulatory compliance checks. | keyword |
 | gen_ai.compliance.violation_code | Code identifying the specific compliance rule that was violated. | keyword |
 | gen_ai.compliance.violation_detected | Indicates if any compliance violation was detected during the interaction. | boolean |
+| gen_ai.guardrail_id | Guardrail ID if a guardrail was executed. | keyword |
 | gen_ai.owasp.description | Description of the OWASP risk triggered. | text |
 | gen_ai.owasp.id | Identifier for the OWASP risk addressed. | keyword |
 | gen_ai.performance.request_size | Size of the request payload in bytes. | long |

diff --git a/packages/aws_bedrock/manifest.yml b/packages/aws_bedrock/manifest.yml
@@ -3,7 +3,7 @@ name: aws_bedrock
 title: Amazon Bedrock
 description: Collect Amazon Bedrock model invocation logs and runtime metrics with Elastic Agent.
 type: integration
-version: "0.12.0"
+version: "0.13.0"
 categories:
   - aws
 conditions: