datadog 2.42.0 → 2.44.0

data/ext/libdatadog_extconf_helpers.rb

@@ -2,6 +2,7 @@
 
 require "rubygems"
 require "pathname"
+require "shellwords"
 
 module Datadog
   # Contains a bunch of shared helpers that get used during building of extensions that link to libdatadog
@@ -10,7 +11,7 @@ module Datadog
   module LibdatadogExtconfHelpers
     # Used to make sure the correct gem version gets loaded, as extconf.rb does not get run with "bundle exec" and thus
     # may see multiple libdatadog versions. See https://github.com/DataDog/dd-trace-rb/pull/2531 for the horror story.
-    LIBDATADOG_VERSION = "~> 40.0.0.2.0"
+    LIBDATADOG_VERSION = "~> 44.0.1.1.0"
 
     # Used as an workaround for a limitation with how dynamic linking works in environments where the datadog gem and
     # libdatadog are moved after the extension gets compiled.
@@ -116,7 +117,10 @@ module Datadog
       extconf_folder:,
       libdatadog_pkgconfig_folder: Libdatadog.pkgconfig_folder,
       gem_dir: Gem.dir,
-      logger: Logging
+      logger: Logging,
+      target_incflags: $INCFLAGS,
+      target_ldflags: $LDFLAGS,
+      target_libs: $libs
     )
       return unless libdatadog_pkgconfig_folder
 
@@ -124,10 +128,9 @@ module Datadog
       libdir = "#{libdatadog_pkgconfig_folder}/../../lib"
       includedir = "#{libdatadog_pkgconfig_folder}/../../include"
 
-      # Set mkmf global variables
-      $INCFLAGS << " -I#{includedir}"
-      $LDFLAGS << " -L#{libdir} -Wl,-rpath,#{libdir}"
-      $libs << " -ldatadog_profiling"
+      target_incflags << " -I#{Shellwords.escape(includedir)}"
+      target_ldflags << " -L#{Shellwords.escape(libdir)} -Wl,-rpath,#{Shellwords.escape(libdir)}"
+      target_libs << " -ldatadog_profiling"
 
       # Add extra relative rpaths using $ORIGIN to handle environments where gems are moved after installation.
       # The excessive escaping is needed to get these special characters through Make and the shell untouched.
@@ -141,10 +144,12 @@ module Datadog
           libdatadog_pkgconfig_folder: libdatadog_pkgconfig_folder,
         ),
       ]
-      extra_relative_rpaths.each { |folder| $LDFLAGS << " -Wl,-rpath,$$$\\\\{ORIGIN\\}/#{folder}" }
+      extra_relative_rpaths.each { |folder| target_ldflags << " -Wl,-rpath,$$$\\\\{ORIGIN\\}/#{Shellwords.escape(folder)}" }
 
-      logger.message("linking with libdatadog (include=#{includedir}, lib=#{libdir})\n")
-      logger.message("[datadog] $LDFLAGS were set to: #{$LDFLAGS.inspect}\n")
+      # includedir/libdir/ldflags might have a `%` in there and `logger.message` bottoms out on `printf` so we need to
+      # use the `message("%s", ...)` format otherwise the logger will complain
+      logger.message("%s", "linking with libdatadog (include=#{includedir}, lib=#{libdir})\n")
+      logger.message("%s", "[datadog] $LDFLAGS were set to: #{target_ldflags.inspect}\n")
 
       true
     end

data/lib/datadog/ai_guard/component.rb

@@ -1,37 +1,46 @@
 # frozen_string_literal: true
 
-require_relative "api_client"
+require "json"
+
+require_relative "ext"
+require_relative "http_client"
 require_relative "evaluation"
 require_relative "evaluation/request"
+require_relative "evaluation/response"
+require_relative "evaluation/client"
 require_relative "evaluation/result"
+require_relative "evaluation/outcome"
 require_relative "evaluation/no_op_result"
 require_relative "evaluation/message"
 require_relative "evaluation/tool_call"
 require_relative "evaluation/content_part"
-require_relative "evaluation/content_builder"
-require_relative "ext"
+require_relative "evaluation/message_builder"
+require_relative "redaction"
+require_relative "redaction/result"
+require_relative "redaction/replacements"
+require_relative "metrics/telemetry"
 
 module Datadog
   module AIGuard
     # Component for API Guard product
     class Component
-      attr_reader :api_client, :logger, :telemetry
+      attr_reader :http_client, :logger, :telemetry
 
       def self.build(settings, logger:, telemetry:)
         return unless settings.respond_to?(:ai_guard) && settings.ai_guard.enabled
 
-        api_client = APIClient.new(
+        http_client = HTTPClient.new(
           endpoint: settings.ai_guard.endpoint,
           api_key: settings.api_key,
           application_key: settings.ai_guard.app_key,
           timeout: settings.ai_guard.timeout_ms / 1_000
         )
 
-        new(api_client, logger: logger, telemetry: telemetry)
+        new(http_client, logger: logger, telemetry: telemetry)
       end
 
-      def initialize(api_client, logger:, telemetry:)
-        @api_client = api_client
+      def initialize(http_client, logger:, telemetry:)
+        @http_client = http_client
         @logger = logger
         @telemetry = telemetry
       end

data/lib/datadog/ai_guard/configuration/ext.rb

@@ -6,6 +6,8 @@ module Datadog
       # This module contains constants for AI Guard component
       module Ext
         ENV_AI_GUARD_ENABLED = "DD_AI_GUARD_ENABLED"
+        ENV_AI_GUARD_REDACTION_ENABLED = "DD_AI_GUARD_REDACTION_ENABLED"
+
         ENV_AI_GUARD_ENDPOINT = "DD_AI_GUARD_ENDPOINT"
         ENV_AI_GUARD_TIMEOUT = "DD_AI_GUARD_TIMEOUT"
         ENV_AI_GUARD_MAX_CONTENT_SIZE = "DD_AI_GUARD_MAX_CONTENT_SIZE"

data/lib/datadog/ai_guard/configuration.rb

@@ -34,6 +34,19 @@ module Datadog
                 o.default false
               end
 
+              # Redact sensitive data from evaluated messages.
+              #
+              # When enabled, values flagged by AI Guard are replaced before the messages are
+              # stored in the span or returned to the caller. Disable to skip redaction entirely.
+              #
+              # @default `DD_AI_GUARD_REDACTION_ENABLED`, otherwise `true`
+              # @return [Boolean]
+              option :redaction_enabled do |o|
+                o.type :bool
+                o.env Ext::ENV_AI_GUARD_REDACTION_ENABLED
+                o.default true
+              end
+
               define_method(:instrument) do |integration_name|
                 return unless enabled # steep:ignore
 

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

@@ -4,74 +4,67 @@ module Datadog
   module AIGuard
     module Contrib
       module RubyLLM
-        # module that gets prepended to RubyLLM::Chat
+        # Protects tool calls before RubyLLM executes them locally
+        #
+        # @api private
         module ChatInstrumentation
-          class << self
-            def evaluate!(messages)
-              ai_guard_messages = messages.flat_map do |message|
-                if message.tool_call?
-                  message.tool_calls.map do |tool_call_id, tool_call|
-                    AIGuard.assistant(id: tool_call_id, tool_name: tool_call.name, arguments: tool_call.arguments.to_s)
-                  end
-                elsif message.tool_result?
-                  build_ai_guard_tool(message)
-                else
-                  build_ai_guard_message(message)
-                end
-              end
+          private
 
-              AIGuard.evaluate(*ai_guard_messages)
-            end
-
-            private
+          def provider_completion(usage_recorder:, stream_tracker: nil, &block)
+            @_datadog_ai_guard_evaluate = true
+            super
+          ensure
+            @_datadog_ai_guard_evaluate = false
+          end
 
-            def build_ai_guard_message(message)
-              content = message.content
+          def preprocessed_messages(list = messages)
+            return super unless @_datadog_ai_guard_evaluate
 
-              case content
-              when ::RubyLLM::Content
-                AIGuard.message(role: message.role) do |m|
-                  m.text(content.text.to_s) if content.text
+            adapter = MessageAdapter.new(list)
 
-                  # Calling attachment.for_llm triggers lazy loading of file contents.
-                  # The result is memoized, so providers won't re-read.
-                  content.attachments.each do |attachment|
-                    case attachment.type
-                    when :image
-                      m.image_url(attachment.for_llm)
-                    when :text
-                      m.text(attachment.for_llm)
-                    end
-                    # Skip :pdf, :audio, :video, :unknown — not supported by AIGuard
-                  end
-                end
-              else
-                AIGuard.message(role: message.role, content: content)
-              end
+            begin
+              converted_messages = adapter.to_ai_guard
+            rescue JSON::JSONError
+              Metrics::Telemetry.report_error
+              return super
+            rescue => e
+              AIGuard.telemetry&.report(e, description: "AI Guard: Failed to convert RubyLLM messages")
+              return super
             end
 
-            def build_ai_guard_tool(message)
-              content = message.content
-              # Tools can return Content or Content::Raw objects (e.g. with attachments),
-              # but AIGuard.tool expects a String. Extract text when content is a Content object.
-              case content
-              when ::RubyLLM::Content
-                content = content.text.to_s
-              end
-              AIGuard.tool(tool_call_id: message.tool_call_id, content: content)
+            evaluation = AIGuard.evaluate(*converted_messages)
+
+            begin
+              redacted_messages = adapter.apply_redactions(evaluation.messages)
+            rescue JSON::JSONError => e
+              AIGuard.telemetry&.report(e, description: "AI Guard: Failed to apply RubyLLM redaction")
+              return super
             end
+
+            super(redacted_messages)
           end
 
-          def complete(&block)
-            Datadog::AIGuard::Contrib::RubyLLM::ChatInstrumentation.evaluate!(messages)
+          def execute_pending_tool_calls(response)
+            response_index = messages.index { |message| message == response }
+            return super unless response_index
 
-            super
-          end
+            adapter = MessageAdapter.new(messages)
+            begin
+              converted_messages = adapter.to_ai_guard
+            rescue JSON::JSONError
+              Metrics::Telemetry.report_error
+              return super
+            end
 
-          def handle_tool_calls(response, &block)
-            Datadog::AIGuard::Contrib::RubyLLM::ChatInstrumentation.evaluate!(messages)
+            evaluation = AIGuard.evaluate(*converted_messages)
+            begin
+              redacted_messages = adapter.apply_redactions(evaluation.messages)
+            rescue JSON::JSONError => e
+              AIGuard.telemetry&.report(e, description: "AI Guard: Failed to apply RubyLLM redaction")
+              return super
+            end
 
-            super
+            super(redacted_messages[response_index])
           end
         end
       end

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

@@ -11,7 +11,7 @@ module Datadog
         class Integration
           include Datadog::AIGuard::Contrib::Integration
 
-          MINIMUM_VERSION = Gem::Version.new("1.0.0")
+          MINIMUM_VERSION = Gem::Version.new("2.0.0")
 
           register_as :ruby_llm, auto_patch: false
 

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

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require_relative "redaction_change_set"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module RubyLLM
+        # RubyLLM messages paired with their AI Guard representation
+        #
+        # @api private
+        class MessageAdapter
+          def initialize(messages)
+            @ruby_llm_messages = messages
+          end
+
+          def to_ai_guard
+            @ai_guard_messages ||= @ruby_llm_messages.map { |message| build_ai_guard_message(message) }
+          end
+
+          def apply_redactions(redacted_messages)
+            ai_guard_messages = to_ai_guard
+            return @ruby_llm_messages if redacted_messages == ai_guard_messages
+
+            @ruby_llm_messages.each_with_index.map do |message, index|
+              original_message = ai_guard_messages[index]
+              redacted_message = redacted_messages[index]
+              next message if redacted_message == original_message
+
+              RedactionChangeSet.new(original_message, redacted_message).apply_to(message)
+            end
+          end
+
+          private
+
+          def build_ai_guard_message(message)
+            builder = Evaluation::MessageBuilder.new
+
+            # NOTE: Hash block arguments are misinterpreted by Steep after `nil.to_h`
+            # steep:ignore:start
+            message.tool_calls.to_h.each do |id, tool_call|
+              builder.tool_call(name: tool_call.name, id: id, arguments: tool_call.arguments)
+            end
+            # steep:ignore:end
+
+            if message.attachments.empty?
+              content = message.content
+            else
+              builder.text(message.content) if message.content && !message.content.empty?
+
+              message.attachments.each do |attachment|
+                next if attachment.provider_file?
+
+                case attachment.type
+                when :text
+                  builder.text(attachment.content)
+                when :image
+                  # NOTE: Local images are Base64-encoded here by `#for_llm`
+                  #       and again when RubyLLM builds the provider request
+                  url = attachment.url? ? attachment.source.to_s : attachment.for_llm
+                  builder.image_url(url)
+                end
+              end
+
+              content = builder.content
+            end
+
+            Evaluation::Message.new(
+              role: message.role, content: content, tool_calls: builder.tool_calls, tool_call_id: message.tool_call_id
+            )
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "message_adapter"
 require_relative "chat_instrumentation"
 
 module Datadog

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