datadog 2.30.0 → 2.32.0

data/lib/datadog/ai_guard/component.rb

@@ -7,6 +7,8 @@ require_relative 'evaluation/result'
 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'
 
 module Datadog

data/lib/datadog/ai_guard/configuration.rb

@@ -1,11 +1,114 @@
 # frozen_string_literal: true
 
+require "uri"
+require_relative "configuration/ext"
+
 module Datadog
   module AIGuard
     # Configuration module for AI Guard
     module Configuration
+      # AI Guard specific settings
+      module Settings
+        def self.extended(base)
+          base = base.singleton_class unless base.is_a?(Class)
+          add_settings!(base)
+        end
+
+        def self.add_settings!(base)
+          base.class_eval do
+            # AI Guard specific configurations.
+            # @public_api
+            #
+            # Steep does not update `self` for this `class_eval` block.
+            # @type self: Datadog::Core::Configuration::Base::_DslContext
+            settings :ai_guard do
+              # Enable AI Guard.
+              #
+              # You can use this option to skip calls to AI Guard API without having to remove library as a whole.
+              #
+              # @default `DD_AI_GUARD_ENABLED`, otherwise `false`
+              # @return [Boolean]
+              option :enabled do |o|
+                o.type :bool
+                o.env Ext::ENV_AI_GUARD_ENABLED
+                o.default false
+              end
+
+              define_method(:instrument) do |integration_name|
+                return unless enabled # steep:ignore
+
+                if (registered_integration = Datadog::AIGuard::Contrib::Integration.registry[integration_name])
+                  klass = registered_integration.klass
+                  if klass.loaded? && klass.compatible?
+                    instance = klass.new
+                    instance.patcher.patch unless instance.patcher.patched?
+                  end
+                end
+              end
+
+              # AI Guard API endpoint path.
+              #
+              # @default `DD_AI_GUARD_ENDPOINT`, otherwise `nil`
+              # @return [String, nil]
+              option :endpoint do |o|
+                o.type :string, nilable: true
+                o.env Ext::ENV_AI_GUARD_ENDPOINT
+
+                o.setter do |value|
+                  next unless value
+
+                  uri = URI(value.to_s)
+                  raise ArgumentError, "Please provide an absolute URI that includes a protocol" unless uri.absolute?
+
+                  uri.to_s.delete_suffix("/")
+                end
+              end
+
+              # Datadog Application key.
+              #
+              # @default `DD_APP_KEY` environment variable, otherwise `nil`
+              # @return [String, nil]
+              option :app_key do |o|
+                o.type :string, nilable: true
+                o.env Ext::ENV_APP_KEY
+              end
+
+              # Request timeout in milliseconds.
+              #
+              # @default `DD_AI_GUARD_TIMEOUT`, otherwise 10 000 ms
+              # @return [Integer]
+              option :timeout_ms do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_TIMEOUT
+                o.default 10_000
+              end
+
+              # Maximum content size in bytes.
+              # Content that exceeds the maximum allowed size is truncated before
+              # being stored in the current span context.
+              #
+              # @default `DD_AI_GUARD_MAX_CONTENT_SIZE`, otherwise 524 228 bytes
+              # @return [Integer]
+              option :max_content_size_bytes do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_MAX_CONTENT_SIZE
+                o.default 512 * 1024
+              end
+
+              # Maximum number of messages.
+              # Older messages are omitted once the message limit is reached.
+              #
+              # @default `DD_AI_GUARD_MAX_MESSAGES_LENGTH`, otherwise 16 messages
+              # @return [Integer]
+              option :max_messages_length do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_MAX_MESSAGES_LENGTH
+                o.default 16
+              end
+            end
+          end
+        end
+      end
     end
   end
 end
-
-require_relative "configuration/settings"

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

@@ -14,13 +14,51 @@ module Datadog
                     AIGuard.assistant(id: tool_call_id, tool_name: tool_call.name, arguments: tool_call.arguments.to_s)
                   end
                 elsif message.tool_result?
-                  AIGuard.tool(tool_call_id: message.tool_call_id, content: message.content)
+                  build_ai_guard_tool(message)
                 else
-                  AIGuard.message(role: message.role, content: message.content)
+                  build_ai_guard_message(message)
                 end
               end
 
-              AIGuard.evaluate(*ai_guard_messages, allow_raise: true)
+              AIGuard.evaluate(*ai_guard_messages)
+            end
+
+            private
+
+            def build_ai_guard_message(message)
+              content = message.content
+
+              case content
+              when ::RubyLLM::Content
+                AIGuard.message(role: message.role) do |m|
+                  m.text(content.text.to_s) if content.text
+
+                  # 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
+            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)
             end
           end
 

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

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Builder for collecting content parts inside a message block.
+      #
+      # Used via the block form of {Datadog::AIGuard.message}:
+      #
+      #   Datadog::AIGuard.message(role: :user) do |m|
+      #     m.text("What's in this image?")
+      #     m.image_url("https://example.com/img.png")
+      #   end
+      class ContentBuilder
+        attr_reader :parts
+
+        def initialize
+          @parts = []
+        end
+
+        def text(text)
+          @parts << ContentPart::Text.new(text)
+        end
+
+        def image_url(url)
+          @parts << ContentPart::ImageURL.new(url)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Namespace for content part types used in multi-modal messages.
+      module ContentPart
+        # A text content part.
+        class Text
+          attr_reader :text
+
+          def initialize(text)
+            @text = text
+          end
+
+          def type
+            :text
+          end
+        end
+
+        # An image URL content part. Accepts an absolute URL or a base64 data URI.
+        class ImageURL
+          attr_reader :url
+
+          def initialize(url)
+            @url = url
+          end
+
+          def type
+            :image_url
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -5,12 +5,14 @@ module Datadog
     module Evaluation
       # Class for emulating AI Guard evaluation result when AI Guard is disabled.
       class NoOpResult
-        attr_reader :action, :reason, :tags
+        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
 
         def initialize
           @action = Result::ALLOW_ACTION
           @reason = "AI Guard is disabled"
           @tags = []
+          @sds_findings = []
+          @tag_probabilities = {}
         end
 
         def allow?

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

@@ -44,15 +44,7 @@ module Datadog
         end
 
         def serialize_messages(messages)
-          serialized_messages = []
-
-          messages.each do |message|
-            serialized_messages << serialize_message(message)
-
-            break if serialized_messages.count == Datadog.configuration.ai_guard.max_messages_length
-          end
-
-          serialized_messages
+          messages.map { |message| serialize_message(message) }
         end
 
         def serialize_message(message)
@@ -69,12 +61,25 @@ module Datadog
                 }
               ]
             }
+          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/result.rb

@@ -9,7 +9,7 @@ module Datadog
         DENY_ACTION = "DENY"
         ABORT_ACTION = "ABORT"
 
-        attr_reader :action, :reason, :tags
+        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
 
         def initialize(raw_response)
           attributes = raw_response.fetch("data").fetch("attributes")
@@ -17,7 +17,9 @@ module Datadog
           @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}\""
         end

data/lib/datadog/ai_guard/evaluation.rb

@@ -6,10 +6,17 @@ module Datadog
     # and creating `ai_guard` span with required tags
     module Evaluation
       class << self
-        def perform(messages, allow_raise: false)
+        def perform(messages, allow_raise: true)
           raise ArgumentError, "Messages must not be empty" if messages&.empty?
 
           Tracing.trace(Ext::SPAN_NAME) do |span, trace|
+            trace.keep!
+            trace.set_tag(
+              Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER,
+              Tracing::Sampling::Ext::Decision::AI_GUARD
+            )
+            trace.set_tag(Ext::EVENT_TAG, true)
+
             if (last_message = messages.last)
               if last_message.tool_call
                 span.set_tag(Ext::TARGET_TAG, "tool")
@@ -34,8 +41,10 @@ module Datadog
             span.set_metastruct_tag(
               Ext::METASTRUCT_TAG,
               {
-                messages: truncate_content(request.serialized_messages),
-                attack_categories: result.tags
+                messages: truncate_content(truncate_messages(request.serialized_messages)),
+                attack_categories: result.tags,
+                sds: result.sds_findings,
+                tag_probs: result.tag_probabilities
               }
             )
 
@@ -56,14 +65,35 @@ module Datadog
 
         private
 
+        def truncate_messages(serialized_messages)
+          max_length = Datadog.configuration.ai_guard.max_messages_length
+          serialized_messages.first(max_length)
+        end
+
+        # Truncates content in serialized messages to stay within the configured byte limit.
+        # For multi-modal messages, only text parts are truncated; image URLs are left intact.
         def truncate_content(serialized_messages)
+          max_bytes = Datadog.configuration.ai_guard.max_content_size_bytes
+
           serialized_messages.map do |message| # steep:ignore
             next message unless message[:content]
 
-            {
-              **message,
-              content: message[:content].byteslice(0, Datadog.configuration.ai_guard.max_content_size_bytes)
-            }
+            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)}
+                else
+                  part
+                end
+              end
+
+              {**message, content: serialized_content}
+            else
+              {
+                **message,
+                content: message[:content].byteslice(0, max_bytes)
+              }
+            end
           end
         end
       end

data/lib/datadog/ai_guard/ext.rb

@@ -10,6 +10,7 @@ module Datadog
       ACTION_TAG = "ai_guard.action"
       REASON_TAG = "ai_guard.reason"
       BLOCKED_TAG = "ai_guard.blocked"
+      EVENT_TAG = "ai_guard.event"
       METASTRUCT_TAG = "ai_guard"
     end
   end

data/lib/datadog/ai_guard.rb

@@ -10,7 +10,7 @@ module Datadog
   module AIGuard
     Core::Configuration::Settings.extend(Configuration::Settings)
 
-    # This error is raised when user passes `allow_raise: true` to Evaluation.perform
+    # This error is raised when `allow_raise` is set to true (the default) in Evaluation.perform
     # and AI Guard considers the messages not safe. Intended to be rescued by the user.
     #
     # WARNING: This name must not change, since front-end is using it.
@@ -68,13 +68,14 @@ module Datadog
       #   One or more message objects to be evaluated.
       # @param allow_raise [Boolean]
       #   Whether this method may raise an exception when evaluation result is not ALLOW.
+      #   Defaults to true.
       #
       # @return [Datadog::AIGuard::Evaluation::Result]
       #   The result of AI Guard evaluation.
       # @raise [Datadog::AIGuard::AIGuardAbortError]
       #   If the evaluation results in DENY or ABORT action and `allow_raise` is set to true
       # @public_api
-      def evaluate(*messages, allow_raise: false)
+      def evaluate(*messages, allow_raise: true)
         if enabled?
           Evaluation.perform(messages, allow_raise: allow_raise)
         else
@@ -84,25 +85,42 @@ module Datadog
 
       # Builds a generic evaluation message.
       #
-      # Example:
+      # Accepts either a string content or a block for multi-modal content parts:
       #
       # ```
+      # # String content:
       # Datadog::AIGuard.message(role: :user, content: "Hello, assistant")
+      #
+      # # Multi-modal content with block:
+      # Datadog::AIGuard.message(role: :user) do |m|
+      #   m.text("What's in this image?")
+      #   m.image_url("https://example.com/img.png")
+      # end
       # ```
       #
       # @param role [Symbol]
       #   The role associated with the message.
       #   Must be one of `:assistant`, `:tool`, `:system`, `:developer`, or `:user`.
-      # @param content [String]
-      #   The textual content of the message.
+      # @param content [String, nil]
+      #   The textual content of the message. Cannot be combined with a block.
+      # @yield [builder] A block for building multi-modal content parts.
+      # @yieldparam builder [Datadog::AIGuard::Evaluation::ContentBuilder]
       #
       # @return [Datadog::AIGuard::Evaluation::Message]
       #   A new message instance with the given role and content.
       # @raise [ArgumentError]
-      #   If an invalid role is provided.
+      #   If both content and a block are provided, or if an invalid role is provided.
       # @public_api
-      def message(role:, content:)
-        Evaluation::Message.new(role: role, content: content)
+      def message(role:, content: nil)
+        if block_given?
+          raise ArgumentError, "Cannot pass both content and a block" if content
+
+          builder = Evaluation::ContentBuilder.new
+          yield builder
+          Evaluation::Message.new(role: role, content: builder.parts)
+        else
+          Evaluation::Message.new(role: role, content: content)
+        end
       end
 
       # Builds an assistant message representing a tool call initiated by the model.

data/lib/datadog/appsec/autoload.rb

@@ -7,7 +7,7 @@ if %w[1 true].include?((Datadog::DATADOG_ENV['DD_APPSEC_ENABLED'] || '').downcas
   rescue => e
     Kernel.warn(
       '[datadog] AppSec failed to instrument. No security check will be performed. error: ' \
-      " #{e.class.name} #{e.message}"
+      " #{e.class}: #{e.message}"
     )
   end
 end

data/lib/datadog/appsec/component.rb

@@ -45,10 +45,15 @@ module Datadog
           settings.appsec.instrument(:devise) unless devise_integration.patcher.patched?
 
           security_engine = SecurityEngine::Engine.new(appsec_settings: settings.appsec, telemetry: telemetry)
-          new(security_engine: security_engine, telemetry: telemetry)
-        rescue
-          Datadog.logger.warn('AppSec is disabled, see logged errors above')
-
+          new(security_engine: security_engine)
+        rescue => e
+          Datadog.logger.warn("AppSec is disabled: #{e.class}: #{e.message}; there may be additional logged errors above")
+
+          # Not reporting to telemetry here because some of the rescued exceptions
+          # have already been reported by the code that raised them
+          # (e.g. SecurityEngine::Engine.new reports WAF init failures).
+          # TODO: reconsider whether telemetry reporting belongs here
+          # (single catch-all) or in the downstream code (as it is now).
           nil
         end
 
@@ -70,11 +75,10 @@ module Datadog
         end
       end
 
-      attr_reader :security_engine, :telemetry
+      attr_reader :security_engine
 
-      def initialize(security_engine:, telemetry:)
+      def initialize(security_engine:)
         @security_engine = security_engine
-        @telemetry = telemetry
       end
 
       def reconfigure!