[ruby@appsec-69389-add-ai-guard-sensetive-data-redaction] Enable AI Guard sensitive data redaction

manifests/ruby.yml

@@ -33,7 +33,14 @@ manifest:
         rails72: '>=2.40.0'
         rails61: '>=2.40.0'
         rails80: '>=2.40.0'
-  tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardTelemetryRedacted: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardTelemetryRedacted:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardTelemetryRequests: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardTelemetryTruncated: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags:
@@ -91,9 +98,30 @@ manifest:
         sinatra32: irrelevant
         sinatra41: irrelevant
         uds-sinatra: irrelevant
-  tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_Redaction: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Redaction:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
   tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabled:  # TODO: a lower version might be supported
     - weblog_declaration:
         '*': missing_feature
@@ -104,17 +132,59 @@ manifest:
         rails61: '>=2.41.0'
   tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabled::test_findings_still_reported:
     - weblog_declaration:
-        rails80: missing_feature
-        rails61: missing_feature
-        uds-rails: missing_feature
-        rails52: missing_feature
-        rails72: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabledTelemetry: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionFailSafe: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionInSDKResponse: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionMultiTurnContext: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionOnBlock: missing_feature
-  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionSkipsStructuralFields: missing_feature
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabledTelemetry:
+    - weblog_declaration:
+        '*': missing_feature
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionFailSafe:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionInSDKResponse:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionMultiTurnContext:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionOnBlock:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
+  tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionSkipsStructuralFields:
+    - weblog_declaration:
+        '*': irrelevant
+        uds-rails: v2.44.0-dev
+        rails52: v2.44.0-dev
+        rails72: v2.44.0-dev
+        rails80: v2.44.0-dev
+        rails61: v2.44.0-dev
   tests/ai_guard/test_ai_guard_sdk.py::Test_RootSpanUserKeep:
     - weblog_declaration:
         "*": v2.31.0-dev

utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb

@@ -4,32 +4,7 @@ class AiGuardController < ApplicationController
   skip_before_action :verify_authenticity_token
 
   def evaluate
-    messages = params.to_unsafe_h.fetch(:_json).flat_map do |message_data|
-      if message_data[:tool_calls]
-        message_data[:tool_calls].map do |tool_call_data|
-          Datadog::AIGuard.assistant(
-            id: tool_call_data[:id],
-            tool_name: tool_call_data.dig(:function, :name),
-            arguments: tool_call_data.dig(:function, :arguments)
-          )
-        end
-      elsif message_data[:tool_call_id]
-        Datadog::AIGuard.tool(tool_call_id: message_data[:tool_call_id], content: message_data[:content])
-      elsif message_data[:content].is_a?(Array)
-        Datadog::AIGuard.message(role: message_data[:role]) do |m|
-          message_data[:content].each do |part|
-            case part[:type]
-            when 'text'
-              m.text(part[:text])
-            when 'image_url'
-              m.image_url(part.dig(:image_url, :url))
-            end
-          end
-        end
-      else
-        Datadog::AIGuard.message(role: message_data[:role], content: message_data[:content])
-      end
-    end
+    messages = params.to_unsafe_h.fetch(:_json).flat_map { |message_data| build_messages(message_data) }
 
     user_id = request.headers['X-User-Id']
     session_id = request.headers['X-Session-Id']
@@ -41,11 +16,14 @@ def evaluate
     response_data = {
       action: result.action,
       reason: result.reason,
-      tags: result.tags,
-      is_blocking_enabled: result.blocking_enabled?
+      tags: result.tags
     }
     response_data[:tag_probs] = result.tag_probabilities if result.respond_to?(:tag_probabilities)
     response_data[:sds] = result.sds_findings if result.respond_to?(:sds_findings)
+    response_data[:messages] = result.messages.map(&:to_h) if result.respond_to?(:messages)
+    if result.respond_to?(:redaction_replacements)
+      response_data[:redaction_replacements] = result.redaction_replacements
+    end
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
@@ -55,4 +33,74 @@ def evaluate
   rescue StandardError => e
     render json: { error: e.to_s, type: e.class.name }, status: 500
   end
+
+  private
+
+  def build_messages(message_data)
+    # TODO: Remove ContentBuilder compatibility after dd-trace-rb 2.44.0 is released
+    if Datadog::AIGuard::Evaluation.const_defined?(:MessageBuilder, false)
+      build_message(message_data)
+    else
+      build_content_builder_messages(message_data)
+    end
+  end
+
+  def build_message(message_data)
+    if message_data[:role] == 'tool'
+      return Datadog::AIGuard.tool(
+        tool_call_id: message_data[:tool_call_id], content: message_data[:content]
+      )
+    end
+
+    message_content = message_data[:content]
+    content = message_content unless message_content.is_a?(Array)
+
+    Datadog::AIGuard.message(role: message_data[:role], content: content) do |message|
+      message_data.fetch(:tool_calls, []).each do |tool_call_data|
+        message.tool_call(
+          id: tool_call_data[:id],
+          name: tool_call_data.dig(:function, :name),
+          arguments: tool_call_data.dig(:function, :arguments)
+        )
+      end
+
+      if message_content.is_a?(Array)
+        message_content.each do |part|
+          case part[:type]
+          when 'text'
+            message.text(part[:text])
+          when 'image_url'
+            message.image_url(part.dig(:image_url, :url))
+          end
+        end
+      end
+    end
+  end
+
+  def build_content_builder_messages(message_data)
+    if message_data[:tool_calls]
+      message_data[:tool_calls].map do |tool_call_data|
+        Datadog::AIGuard.assistant(
+          id: tool_call_data[:id],
+          tool_name: tool_call_data.dig(:function, :name),
+          arguments: tool_call_data.dig(:function, :arguments)
+        )
+      end
+    elsif message_data[:tool_call_id]
+      Datadog::AIGuard.tool(tool_call_id: message_data[:tool_call_id], content: message_data[:content])
+    elsif message_data[:content].is_a?(Array)
+      Datadog::AIGuard.message(role: message_data[:role]) do |message|
+        message_data[:content].each do |part|
+          case part[:type]
+          when 'text'
+            message.text(part[:text])
+          when 'image_url'
+            message.image_url(part.dig(:image_url, :url))
+          end
+        end
+      end
+    else
+      Datadog::AIGuard.message(role: message_data[:role], content: message_data[:content])
+    end
+  end
 end

utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb

@@ -4,32 +4,7 @@ class AiGuardController < ApplicationController
   skip_before_action :verify_authenticity_token
 
   def evaluate
-    messages = params.to_unsafe_h.fetch(:_json).flat_map do |message_data|
-      if message_data[:tool_calls]
-        message_data[:tool_calls].map do |tool_call_data|
-          Datadog::AIGuard.assistant(
-            id: tool_call_data[:id],
-            tool_name: tool_call_data.dig(:function, :name),
-            arguments: tool_call_data.dig(:function, :arguments)
-          )
-        end
-      elsif message_data[:tool_call_id]
-        Datadog::AIGuard.tool(tool_call_id: message_data[:tool_call_id], content: message_data[:content])
-      elsif message_data[:content].is_a?(Array)
-        Datadog::AIGuard.message(role: message_data[:role]) do |m|
-          message_data[:content].each do |part|
-            case part[:type]
-            when 'text'
-              m.text(part[:text])
-            when 'image_url'
-              m.image_url(part.dig(:image_url, :url))
-            end
-          end
-        end
-      else
-        Datadog::AIGuard.message(role: message_data[:role], content: message_data[:content])
-      end
-    end
+    messages = params.to_unsafe_h.fetch(:_json).flat_map { |message_data| build_messages(message_data) }
 
     user_id = request.headers['X-User-Id']
     session_id = request.headers['X-Session-Id']
@@ -41,11 +16,14 @@ def evaluate
     response_data = {
       action: result.action,
       reason: result.reason,
-      tags: result.tags,
-      is_blocking_enabled: result.blocking_enabled?
+      tags: result.tags
     }
     response_data[:tag_probs] = result.tag_probabilities if result.respond_to?(:tag_probabilities)
     response_data[:sds] = result.sds_findings if result.respond_to?(:sds_findings)
+    response_data[:messages] = result.messages.map(&:to_h) if result.respond_to?(:messages)
+    if result.respond_to?(:redaction_replacements)
+      response_data[:redaction_replacements] = result.redaction_replacements
+    end
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
@@ -55,4 +33,74 @@ def evaluate
   rescue StandardError => e
     render json: { error: e.to_s, type: e.class.name }, status: 500
   end
+
+  private
+
+  def build_messages(message_data)
+    # TODO: Remove ContentBuilder compatibility after dd-trace-rb 2.44.0 is released
+    if Datadog::AIGuard::Evaluation.const_defined?(:MessageBuilder, false)
+      build_message(message_data)
+    else
+      build_content_builder_messages(message_data)
+    end
+  end
+
+  def build_message(message_data)
+    if message_data[:role] == 'tool'
+      return Datadog::AIGuard.tool(
+        tool_call_id: message_data[:tool_call_id], content: message_data[:content]
+      )
+    end
+
+    message_content = message_data[:content]
+    content = message_content unless message_content.is_a?(Array)
+
+    Datadog::AIGuard.message(role: message_data[:role], content: content) do |message|
+      message_data.fetch(:tool_calls, []).each do |tool_call_data|
+        message.tool_call(
+          id: tool_call_data[:id],
+          name: tool_call_data.dig(:function, :name),
+          arguments: tool_call_data.dig(:function, :arguments)
+        )
+      end
+
+      if message_content.is_a?(Array)
+        message_content.each do |part|
+          case part[:type]
+          when 'text'
+            message.text(part[:text])
+          when 'image_url'
+            message.image_url(part.dig(:image_url, :url))
+          end
+        end
+      end
+    end
+  end
+
+  def build_content_builder_messages(message_data)
+    if message_data[:tool_calls]
+      message_data[:tool_calls].map do |tool_call_data|
+        Datadog::AIGuard.assistant(
+          id: tool_call_data[:id],
+          tool_name: tool_call_data.dig(:function, :name),
+          arguments: tool_call_data.dig(:function, :arguments)
+        )
+      end
+    elsif message_data[:tool_call_id]
+      Datadog::AIGuard.tool(tool_call_id: message_data[:tool_call_id], content: message_data[:content])
+    elsif message_data[:content].is_a?(Array)
+      Datadog::AIGuard.message(role: message_data[:role]) do |message|
+        message_data[:content].each do |part|
+          case part[:type]
+          when 'text'
+            message.text(part[:text])
+          when 'image_url'
+            message.image_url(part.dig(:image_url, :url))
+          end
+        end
+      end
+    else
+      Datadog::AIGuard.message(role: message_data[:role], content: message_data[:content])
+    end
+  end
 end