datadog 2.39.0 → 2.47.0

data/lib/datadog/ai_guard/evaluation/no_op_result.rb

@@ -5,14 +5,18 @@ module Datadog
     module Evaluation
       # Class for emulating AI Guard evaluation result when AI Guard is disabled.
       class NoOpResult
-        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
+        attr_reader :action, :reason, :tags, :sds_findings, :redaction_replacements,
+          :tag_probabilities, :messages
 
-        def initialize
+        def initialize(messages)
           @action = Result::ALLOW_ACTION
           @reason = "AI Guard is disabled"
+
           @tags = []
+          @messages = messages
           @sds_findings = []
           @tag_probabilities = {}
+          @redaction_replacements = []
         end
 
         def allow?
@@ -26,10 +30,6 @@ module Datadog
         def abort?
           false
         end
-
-        def blocking_enabled?
-          false
-        end
       end
     end
   end

data/lib/datadog/ai_guard/evaluation/outcome.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Internal outcome of an evaluation request
+      #
+      # @api private
+      class Outcome
+        attr_reader :result, :redaction
+
+        def initialize(result:, redaction:, blocking_enabled:)
+          @result = result
+          @redaction = redaction
+          @blocking_enabled = blocking_enabled
+        end
+
+        def block?
+          @blocking_enabled && (result.deny? || result.abort?)
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/evaluation/request.rb

@@ -3,83 +3,27 @@
 module Datadog
   module AIGuard
     module Evaluation
-      # Request builds the request body from an array of messages and processes the response
+      # Request builds the request body for the AI Guard /evaluate endpoint
       class Request
         REQUEST_PATH = "/evaluate"
 
-        attr_reader :serialized_messages
-
         def initialize(messages)
-          @serialized_messages = serialize_messages(messages)
+          @messages = messages
         end
 
-        def perform
-          api_client = AIGuard.api_client
-
-          # This should never happen, as we are only calling this method when AI Guard is enabled,
-          # and this means the API Client was not initialized properly.
-          #
-          # Please report this at https://github.com/datadog/dd-trace-rb/blob/master/CONTRIBUTING.md#found-a-bug
-          raise "AI Guard API Client not initialized" unless api_client
-
-          raw_response = api_client.post(REQUEST_PATH, body: build_request_body)
-
-          Result.new(raw_response)
-        end
-
-        private
-
-        def build_request_body
+        def body
           {
             data: {
               attributes: {
-                messages: @serialized_messages,
+                messages: @messages.map(&:to_h),
                 meta: {
                   service: Datadog.configuration.service,
-                  env: Datadog.configuration.env
-                }
-              }
-            }
+                  env: Datadog.configuration.env,
+                },
+              },
+            },
           }
         end
-
-        def serialize_messages(messages)
-          messages.map { |message| serialize_message(message) }
-        end
-
-        def serialize_message(message)
-          if message.tool_call
-            {
-              role: message.role,
-              tool_calls: [
-                {
-                  id: message.tool_call.id,
-                  function: {
-                    name: message.tool_call.tool_name,
-                    arguments: message.tool_call.arguments
-                  }
-                }
-              ]
-            }
-          elsif message.content.is_a?(::Array)
-            {role: message.role, content: serialize_content_parts(message.content)}
-          elsif message.tool_call_id
-            {role: message.role, tool_call_id: message.tool_call_id, content: message.content}
-          else
-            {role: message.role, content: message.content}
-          end
-        end
-
-        def serialize_content_parts(parts)
-          parts.map do |part|
-            case part
-            when ContentPart::Text
-              {type: "text", text: part.text}
-            when ContentPart::ImageURL
-              {type: "image_url", image_url: {url: part.url}}
-            end
-          end
-        end
       end
     end
   end

data/lib/datadog/ai_guard/evaluation/response.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Wrapper class for evaluation API response
+      #
+      # @api private
+      class Response
+        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities,
+          :redaction_replacements
+
+        def initialize(raw_response)
+          attributes = raw_response.fetch("data").fetch("attributes")
+
+          @action = attributes.fetch("action")
+          @reason = attributes.fetch("reason")
+          @tags = attributes.fetch("tags")
+          @tag_probabilities = attributes.fetch("tag_probs")
+          @is_blocking_enabled = attributes.fetch("is_blocking_enabled")
+          @sds_findings = attributes.fetch("sds_findings", [])
+          @redaction_replacements = attributes.fetch("redaction_replacements", [])
+        rescue KeyError => e
+          raise AIGuardClientError, "Missing key: \"#{e.key}\""
+        end
+
+        def blocking_enabled?
+          !!@is_blocking_enabled
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/evaluation/result.rb

@@ -3,25 +3,22 @@
 module Datadog
   module AIGuard
     module Evaluation
-      # Wrapper class for evaluation API response
       class Result
         ALLOW_ACTION = "ALLOW"
         DENY_ACTION = "DENY"
         ABORT_ACTION = "ABORT"
 
-        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
-
-        def initialize(raw_response)
-          attributes = raw_response.fetch("data").fetch("attributes")
-
-          @action = attributes.fetch("action")
-          @reason = attributes.fetch("reason")
-          @tags = attributes.fetch("tags")
-          @tag_probabilities = attributes.fetch("tag_probs")
-          @is_blocking_enabled = attributes.fetch("is_blocking_enabled")
-          @sds_findings = attributes.fetch("sds_findings", [])
-        rescue KeyError => e
-          raise AIGuardClientError, "Missing key: \"#{e.key}\""
+        attr_reader :messages, :action, :reason, :tags, :sds_findings,
+          :tag_probabilities, :redaction_replacements
+
+        def initialize(messages, action:, reason:, tags:, sds_findings:, tag_probabilities:, redaction_replacements:)
+          @messages = messages
+          @action = action
+          @reason = reason
+          @tags = tags
+          @sds_findings = sds_findings
+          @tag_probabilities = tag_probabilities
+          @redaction_replacements = redaction_replacements
         end
 
         def allow?
@@ -35,10 +32,6 @@ module Datadog
         def abort?
           action == ABORT_ACTION
         end
-
-        def blocking_enabled?
-          !!@is_blocking_enabled
-        end
       end
     end
   end

data/lib/datadog/ai_guard/evaluation/tool_call.rb

@@ -3,14 +3,27 @@
 module Datadog
   module AIGuard
     module Evaluation
-      # Tool call class for AI Guard
+      # Represents a tool invocation requested by an assistant message
       class ToolCall
         attr_reader :tool_name, :id, :arguments
 
         def initialize(tool_name, id:, arguments:)
           @tool_name = tool_name
           @id = id
-          @arguments = arguments
+          @arguments =
+            case arguments
+            when String then arguments
+            when Hash then JSON.generate(arguments)
+            else raise ArgumentError, "Tool call arguments must be a String or Hash"
+            end
+        end
+
+        def copy(arguments:)
+          ToolCall.new(tool_name, id: id, arguments: arguments)
+        end
+
+        def to_h
+          {id: id, function: {name: tool_name, arguments: arguments}}
         end
       end
     end

data/lib/datadog/ai_guard/evaluation.rb

@@ -15,6 +15,7 @@ module Datadog
               Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER,
               Tracing::Sampling::Ext::Decision::AI_GUARD
             )
+            trace.set_distributed_source(Ext::PRODUCT_BIT)
             trace.set_tag(Ext::EVENT_TAG, true)
             trace.set_tag(Ext::TRACE_EXECUTED_TAG, "1")
 
@@ -24,37 +25,51 @@ module Datadog
             end
 
             if (last_message = messages.last)
-              if last_message.tool_call
+              if (tool_call = last_message.tool_calls.first)
                 span.set_tag(Ext::TARGET_TAG, "tool")
-                span.set_tag(Ext::TOOL_NAME_TAG, last_message.tool_call.tool_name)
+                span.set_tag(Ext::TOOL_NAME_TAG, tool_call.tool_name)
               elsif last_message.tool_call_id
                 span.set_tag(Ext::TARGET_TAG, "tool")
 
-                if (tool_call_message = messages.find { |m| m.tool_call&.id == last_message.tool_call_id })
-                  span.set_tag(Ext::TOOL_NAME_TAG, tool_call_message.tool_call.tool_name) # steep:ignore
+                messages.reverse_each do |message|
+                  # @type var tool_call: ToolCall?
+                  tool_call = message.tool_calls.find { |call| call.id == last_message.tool_call_id }
+                  break span.set_tag(Ext::TOOL_NAME_TAG, tool_call.tool_name) if tool_call
                 end
               else
                 span.set_tag(Ext::TARGET_TAG, "prompt")
               end
             end
 
-            request = Request.new(messages)
-            result = request.perform
+            outcome =
+              begin
+                Client.evaluate(messages)
+              rescue
+                Metrics::Telemetry.report_error
+                raise
+              end
+
+            result = outcome.result
+            redaction = outcome.redaction
 
             span.set_tag(Ext::ACTION_TAG, result.action)
             span.set_tag(Ext::REASON_TAG, result.reason)
+            span.set_tag(Ext::REDACTED_TAG, redaction.redacted?) if redaction.performed?
 
             span.set_metastruct_tag(
               Ext::METASTRUCT_TAG,
               {
-                messages: truncate_content(truncate_messages(request.serialized_messages)),
+                messages: truncate_content(truncate_messages(result.messages).map(&:to_h)),
                 attack_categories: result.tags,
                 sds: result.sds_findings,
-                tag_probs: result.tag_probabilities
+                tag_probs: result.tag_probabilities,
               }
             )
 
-            if allow_raise && (result.deny? || result.abort?) && result.blocking_enabled?
+            blocked = allow_raise && outcome.block?
+            Metrics::Telemetry.report_evaluation(outcome, blocked: blocked)
+
+            if blocked
               span.set_tag(Ext::BLOCKED_TAG, true)
               raise AIGuardAbortError.new(action: result.action, reason: result.reason, tags: result.tags)
             end
@@ -63,10 +78,10 @@ module Datadog
           end
         end
 
-        def perform_no_op
+        def perform_no_op(messages)
           AIGuard.logger&.warn("AI Guard is disabled, messages were not evaluated")
 
-          NoOpResult.new
+          NoOpResult.new(messages)
         end
 
         private
@@ -81,13 +96,13 @@ module Datadog
         def truncate_content(serialized_messages)
           max_bytes = Datadog.configuration.ai_guard.max_content_size_bytes
 
-          serialized_messages.map do |message| # steep:ignore
+          serialized_messages.map do |message|
             next message unless message[:content]
 
             if message[:content].is_a?(::Array)
               serialized_content = message[:content].map do |part|
                 if part[:text]
-                  {**part, text: part[:text].to_s.byteslice(0, max_bytes)}
+                  {**part, text: truncate_content_value(part[:text].to_s, max_bytes)}
                 else
                   part
                 end
@@ -97,11 +112,18 @@ module Datadog
             else
               {
                 **message,
-                content: message[:content].byteslice(0, max_bytes)
+                content: truncate_content_value(message[:content], max_bytes),
               }
             end
           end
         end
+
+        def truncate_content_value(content, max_bytes)
+          truncated = content.byteslice(0, max_bytes).to_s
+          return truncated if truncated.valid_encoding?
+
+          truncated.scrub("")
+        end
       end
     end
   end

data/lib/datadog/ai_guard/ext.rb

@@ -4,12 +4,15 @@ module Datadog
   module AIGuard
     # AI Guard specific constants
     module Ext
+      PRODUCT_BIT = 0x20
+
       SPAN_NAME = "ai_guard"
       TARGET_TAG = "ai_guard.target"
       TOOL_NAME_TAG = "ai_guard.tool_name"
       ACTION_TAG = "ai_guard.action"
       REASON_TAG = "ai_guard.reason"
       BLOCKED_TAG = "ai_guard.blocked"
+      REDACTED_TAG = "ai_guard.redacted"
       EVENT_TAG = "ai_guard.event"
 
       METASTRUCT_TAG = "ai_guard"
@@ -22,7 +25,7 @@ module Datadog
       TRACE_ANOMALY_DETECTION_TAGS = [
         TRACE_HTTP_USERAGENT_TAG,
         TRACE_HTTP_CLIENT_IP_TAG,
-        TRACE_NETWORK_CLIENT_IP_TAG
+        TRACE_NETWORK_CLIENT_IP_TAG,
       ].freeze
     end
   end

data/lib/datadog/ai_guard/{api_client.rb → http_client.rb}

@@ -1,14 +1,16 @@
 # frozen_string_literal: true
 
 require "uri"
-require "net/http"
 require "json"
+require "net/http"
 
 module Datadog
   module AIGuard
-    # API Client for AI Guard API.
-    # Uses net/http to perform request. Raises on client and server errors.
-    class APIClient
+    # HTTP client for the AI Guard API
+    # Uses net/http to perform requests, raises on client and server errors
+    #
+    # @api private
+    class HTTPClient
       DEFAULT_SITE = "app.datadoghq.com"
       DEFAULT_PATH = "/api/v2/ai-guard"
 
@@ -32,7 +34,7 @@ module Datadog
           "DD-AI-GUARD-VERSION": Datadog::VERSION::STRING,
           "DD-AI-GUARD-SOURCE": "SDK",
           "DD-AI-GUARD-LANGUAGE": "ruby",
-          "content-type": "application/json"
+          "content-type": "application/json",
         }.freeze
       end
 

data/lib/datadog/ai_guard/metrics/telemetry.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    # Namespace for metrics
+    #
+    # @api private
+    module Metrics
+      # A module for telemetry metrics reporting
+      module Telemetry
+        NAMESPACE = "ai_guard"
+        CLIENT_ERROR = "client_error"
+        REDACTION_ERROR = "redaction_error"
+
+        module_function
+
+        def report_evaluation(outcome, blocked:)
+          telemetry = AIGuard.telemetry
+          return unless telemetry
+
+          result = outcome.result
+          redaction = outcome.redaction
+
+          tags = {action: result.action, block: blocked.to_s, error: "false"}
+          tags[:redacted] = redaction.redacted?.to_s if redaction.performed?
+
+          telemetry.inc(NAMESPACE, "requests", 1, tags: tags)
+          if redaction.failures_count.positive?
+            telemetry.inc(NAMESPACE, "error", redaction.failures_count, tags: {type: REDACTION_ERROR})
+          end
+        end
+
+        def report_error
+          telemetry = AIGuard.telemetry
+          return unless telemetry
+
+          telemetry.inc(NAMESPACE, "requests", 1, tags: {error: "true"})
+          telemetry.inc(NAMESPACE, "error", 1, tags: {type: CLIENT_ERROR})
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/redaction/replacements.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require "forwardable"
+
+module Datadog
+  module AIGuard
+    module Redaction
+      # Converts raw backend redaction replacements into a conflict-free
+      # path-to-replacement map
+      #
+      # @api private
+      class Replacements
+        extend Forwardable
+
+        TEXT_PATH_PATTERN = /\Amessages\[([0-9]+)\]\.content\[([0-9]+)\]\.text\z/
+        CONTENT_PATH_PATTERN = /\Amessages\[([0-9]+)\]\.content\z/
+        ARGUMENTS_PATH_PATTERN = /\Amessages\[([0-9]+)\]\.tool_calls\[([0-9]+)\]\.function\.arguments\z/
+
+        def_delegators :@replacements, :each, :empty?
+
+        attr_reader :failures_count
+
+        def initialize(raw_replacements)
+          @failures_count = 0
+          @replacements = build(raw_replacements)
+        end
+
+        private
+
+        # NOTE: Identical duplicates collapse into a single entry,
+        #       conflicting replacements remove the path and record one failure,
+        #       and malformed or unsupported entries are skipped and counted individually
+        #
+        # Examples
+        #
+        #   translates this
+        #
+        #   [
+        #     {"path" => "messages[0].content", "replacement" => "<redacted>"},
+        #     {"path" => "messages[1].content[2].text", "replacement" => "<redacted>"},
+        #     {"path" => "messages[2].tool_calls[0].function.arguments", "replacement" => "{}"}
+        #   ]
+        #
+        #   into this
+        #
+        #   {
+        #     [0, :content] => "<redacted>",
+        #     [1, :text, 2] => "<redacted>",
+        #     [2, :arguments, 0] => "{}"
+        #   }
+        def build(raw_replacements)
+          unless raw_replacements.is_a?(::Array)
+            @failures_count += 1
+            return {}
+          end
+
+          conflicting = {}
+          raw_replacements.each_with_object({}) do |entry, replacements|
+            next @failures_count += 1 unless entry.is_a?(::Hash)
+
+            raw_path = entry["path"]
+            replacement = entry["replacement"]
+
+            if !raw_path.is_a?(::String) || raw_path.empty? || !replacement.is_a?(::String)
+              next @failures_count += 1
+            end
+
+            # @type var path: Replacements::path?
+            path =
+              case raw_path
+              when CONTENT_PATH_PATTERN
+                [Regexp.last_match(1).to_i, :content]
+              when TEXT_PATH_PATTERN
+                [Regexp.last_match(1).to_i, :text, Regexp.last_match(2).to_i]
+              when ARGUMENTS_PATH_PATTERN
+                [Regexp.last_match(1).to_i, :arguments, Regexp.last_match(2).to_i]
+              end
+
+            next @failures_count += 1 unless path
+            next if conflicting.key?(path)
+
+            if replacements.key?(path)
+              next if replacements[path] == replacement
+
+              replacements.delete(path)
+              conflicting[path] = true
+
+              next @failures_count += 1
+            end
+
+            replacements[path] = replacement
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/redaction/result.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Redaction
+      # A result of applying redaction replacements
+      #
+      # @api private
+      class Result
+        attr_reader :messages, :applied_count, :failures_count
+
+        def initialize(messages, applied_count:, failures_count:, performed: true)
+          @messages = messages
+          @applied_count = applied_count
+          @failures_count = failures_count
+          @performed = performed
+        end
+
+        def performed?
+          @performed
+        end
+
+        def redacted?
+          applied_count.positive?
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/redaction.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    # Sensitive data redaction for evaluation messages
+    #
+    # @api private
+    module Redaction
+      class << self
+        def skip(messages)
+          Redaction::Result.new(messages, applied_count: 0, failures_count: 0, performed: false)
+        end
+
+        def perform(messages, replacements:)
+          redaction_replacements = Replacements.new(replacements)
+          failures_count = redaction_replacements.failures_count
+
+          if redaction_replacements.empty?
+            return Redaction::Result.new(messages, applied_count: 0, failures_count: failures_count)
+          end
+
+          applied_count = 0
+          redacted_messages = messages.dup
+
+          redaction_replacements.each do |path, replacement|
+            index = path[0]
+            redacted_message = redact(redacted_messages[index], path: path, replacement: replacement)
+
+            next failures_count += 1 unless redacted_message
+
+            redacted_messages[index] = redacted_message
+            applied_count += 1
+          rescue
+            failures_count += 1
+          end
+
+          Redaction::Result.new(
+            redacted_messages, applied_count: applied_count, failures_count: failures_count
+          )
+        end
+
+        private
+
+        def redact(message, path:, replacement:)
+          return unless message.is_a?(Evaluation::Message)
+
+          _, kind, index = path
+
+          case kind
+          when :content
+            return unless message.content.is_a?(::String)
+
+            message.copy(content: replacement)
+          when :text
+            # @type var index: Integer
+            content = message.content
+            return unless content.is_a?(::Array)
+
+            part = content[index]
+            return if !part.is_a?(Evaluation::ContentPart::Text) || !part.text.is_a?(::String)
+
+            redacted_content = ::Array.new(content)
+            redacted_content[index] = part.copy(text: replacement)
+
+            message.copy(content: redacted_content)
+          when :arguments
+            # @type var index: Integer
+            tool_call = message.tool_calls[index]
+            return unless tool_call && tool_call.arguments.is_a?(::String)
+
+            redacted_tool_calls = ::Array.new(message.tool_calls)
+            redacted_tool_calls[index] = tool_call.copy(arguments: replacement)
+
+            message.copy(tool_calls: redacted_tool_calls)
+          end
+        end
+      end
+    end
+  end
+end