datadog 2.43.0 → 2.44.0

data/lib/datadog/ai_guard/contrib/ruby_llm/redaction_change_set.rb

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+require "json"
+require "stringio"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module RubyLLM
+        # Changes required to redact a RubyLLM message
+        #
+        # @api private
+        class RedactionChangeSet
+          def initialize(original_message, redacted_message)
+            @content_replacement = content_replacement(original_message.content, redacted_message.content)
+            @content_part_replacements = content_part_replacements(original_message.content, redacted_message.content)
+            @tool_call_replacements = tool_call_replacements(original_message.tool_calls, redacted_message.tool_calls)
+          end
+
+          def apply_to(message)
+            return message unless changed?
+
+            # NOTE: Preserve `raw_content` after provider-specific redaction is implemented.
+            #       Unchanged raw content could bypass redaction
+            if message.raw_content
+              Datadog.logger.warn("AI Guard omitted RubyLLM raw content because it may contain sensitive data")
+            end
+
+            content = build_content(message)
+            attachments = build_attachments(message)
+            tool_calls = build_tool_calls(message)
+
+            changed_message = ::RubyLLM::Message.new(
+              content: content, attachments: attachments, tool_calls: tool_calls,
+              role: message.role, model: message.model, tool_call_id: message.tool_call_id,
+              raw: message.raw, thinking: message.thinking, citations: message.citations,
+              finish_reason: message.finish_reason, server_tool_calls: message.server_tool_calls,
+              raw_reasoning: message.raw_reasoning, cache_until_here: message.cache_until_here?
+            )
+            changed_message.conversation = message.conversation
+            changed_message
+          end
+
+          private
+
+          # NOTE: Steep treats original and redacted content shapes as independent
+          def content_replacement(original_content, redacted_content) # steep:ignore MethodBodyTypeMismatch
+            return if original_content.is_a?(Array) || redacted_content == original_content
+
+            redacted_content
+          end
+
+          def content_part_replacements(original_content, redacted_content)
+            return {} unless original_content.is_a?(Array)
+
+            replacements = {}
+            original_content.each_with_index do |content_part, index|
+              next if content_part == redacted_content[index] # steep:ignore NoMethod
+
+              replacements[index] = redacted_content[index] # steep:ignore NoMethod
+            end
+
+            replacements
+          end
+
+          def tool_call_replacements(original_tool_calls, redacted_tool_calls)
+            replacements = {}
+            original_tool_calls.each_with_index do |tool_call, index|
+              next if redacted_tool_calls[index] == tool_call
+
+              replacements[tool_call.id] = redacted_tool_calls[index]
+            end
+
+            replacements
+          end
+
+          def changed?
+            !!@content_replacement || !@content_part_replacements.empty? ||
+              !@tool_call_replacements.empty?
+          end
+
+          def build_content(message)
+            return @content_replacement if @content_replacement
+            return message.content if message.content.to_s.empty?
+
+            replacement = @content_part_replacements[0]
+            replacement ? replacement.text : message.content
+          end
+
+          def build_attachments(message)
+            return message.attachments if @content_part_replacements.empty?
+
+            content_part_index = message.content.to_s.empty? ? 0 : 1
+            # @type var attachments: Array[::RubyLLM::Attachment]?
+            attachments = nil
+
+            message.attachments.each_with_index do |attachment, attachment_index|
+              next if attachment.provider_file?
+              next unless attachment.type == :text || attachment.type == :image
+
+              replacement = @content_part_replacements[content_part_index]
+              content_part_index += 1
+
+              next unless replacement && attachment.type == :text
+
+              attachments ||= message.attachments.dup
+              attachments[attachment_index] = ::RubyLLM::Attachment.new( # steep:ignore NoMethod
+                StringIO.new(replacement.text), filename: attachment.filename, config: attachment.config
+              )
+            end
+
+            attachments || message.attachments
+          end
+
+          def build_tool_calls(message)
+            return message.tool_calls if @tool_call_replacements.empty?
+
+            # @type var tool_calls: Hash[String, ::RubyLLM::ToolCall]
+            tool_calls = message.tool_calls.dup # steep:ignore IncompatibleAssignment
+            @tool_call_replacements.each do |id, replacement|
+              tool_call = tool_calls.fetch(id)
+              tool_calls[id] = ::RubyLLM::ToolCall.new(
+                id: tool_call.id, name: tool_call.name, remote: tool_call.remote?,
+                thought_signature: tool_call.thought_signature,
+                arguments: JSON.parse(replacement.arguments)
+              )
+            end
+
+            tool_calls
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Client for evaluation API requests
+      #
+      # @api private
+      module Client
+        def self.evaluate(messages)
+          request = Request.new(messages)
+          http_client = AIGuard.http_client
+
+          # This should never happen, as we are only calling this method when AI Guard is enabled,
+          # and this means the HTTP 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 HTTP client not initialized" unless http_client
+
+          response = Response.new(
+            http_client.post(Request::REQUEST_PATH, body: request.body)
+          )
+
+          redaction =
+            if Datadog.configuration.ai_guard.redaction_enabled
+              Redaction.perform(messages, replacements: response.redaction_replacements)
+            else
+              Redaction.skip(messages)
+            end
+
+          result = Result.new(
+            redaction.messages,
+            action: response.action,
+            reason: response.reason,
+            tags: response.tags,
+            sds_findings: response.sds_findings,
+            tag_probabilities: response.tag_probabilities,
+            redaction_replacements: response.redaction_replacements
+          )
+
+          Outcome.new(
+            result: result,
+            redaction: redaction,
+            blocking_enabled: response.blocking_enabled?
+          )
+        end
+      end
+    end
+  end
+end

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

@@ -16,6 +16,14 @@ module Datadog
           def type
             :text
           end
+
+          def copy(text:)
+            Text.new(text)
+          end
+
+          def to_h
+            {type: type.to_s, text: text}
+          end
         end
 
         # An image URL content part. Accepts an absolute URL or a base64 data URI.
@@ -29,6 +37,10 @@ module Datadog
           def type
             :image_url
           end
+
+          def to_h
+            {type: type.to_s, image_url: {url: url}}
+          end
         end
       end
     end

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

@@ -3,21 +3,39 @@
 module Datadog
   module AIGuard
     module Evaluation
-      # Message class for AI Guard
+      # Represents a message submitted for evaluation
       class Message
-        attr_reader :role, :content, :tool_call, :tool_call_id
+        attr_reader :role, :content, :tool_calls, :tool_call_id
 
-        def initialize(role:, content: nil, tool_call: nil, tool_call_id: nil)
+        def initialize(role:, content: nil, tool_calls: [], tool_call_id: nil)
           raise ArgumentError, "Role must be set to a non-empty value" if role.to_s.empty?
+          raise ArgumentError, "Tool calls must be an Array" unless tool_calls.is_a?(::Array)
+
+          if tool_calls.any? { |tool_call| !tool_call.is_a?(ToolCall) }
+            raise ArgumentError, "Tool calls must contain only #{ToolCall.name} instances"
+          end
 
           @role = role.to_sym
           @content = content
-          @tool_call = tool_call
+          @tool_calls = tool_calls
           @tool_call_id = tool_call_id
+        end
 
-          if @tool_call && !@tool_call.is_a?(ToolCall)
-            raise ArgumentError, "Expected an instance of #{ToolCall.name} for :tool_call argument"
-          end
+        def copy(content: self.content, tool_calls: self.tool_calls)
+          Message.new(role: role, content: content, tool_calls: tool_calls, tool_call_id: tool_call_id)
+        end
+
+        def to_h
+          serialized = {
+            role: role,
+            content: content.is_a?(::Array) ? content.map(&:to_h) : content,
+            tool_call_id: tool_call_id,
+          }
+
+          serialized[:tool_calls] = tool_calls.map(&:to_h) unless tool_calls.empty?
+          serialized.compact!
+
+          serialized
         end
       end
     end

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

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Builds a message from content parts and tool calls
+      #
+      # @public_api
+      class MessageBuilder
+        attr_reader :content, :tool_calls
+
+        def initialize
+          @content = []
+          @tool_calls = []
+        end
+
+        def text(text)
+          @content << ContentPart::Text.new(text)
+        end
+
+        def image_url(url)
+          @content << ContentPart::ImageURL.new(url)
+        end
+
+        def tool_call(name:, id:, arguments:)
+          @tool_calls << ToolCall.new(name, id: id.to_s, arguments: arguments)
+        end
+      end
+    end
+  end
+end

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,37 +3,19 @@
 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,
@@ -42,44 +24,6 @@ module Datadog
             },
           }
         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

@@ -25,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,
               }
             )
 
-            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
@@ -64,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
@@ -82,7 +96,7 @@ 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)

data/lib/datadog/ai_guard/ext.rb

@@ -12,6 +12,7 @@ module Datadog
       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"

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"