datadog 2.21.0 → 2.38.0

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      # Base provides features that are shared across all integrations
+      module Integration
+        RegisteredIntegration = Struct.new(:name, :klass, :options)
+
+        @registry = {}
+
+        # Class-level methods for Integration
+        module ClassMethods
+          def register_as(name, options = {})
+            Integration.register(self, name, options)
+          end
+
+          def compatible?
+            true
+          end
+        end
+
+        def self.included(base)
+          base.extend(ClassMethods)
+        end
+
+        def self.register(integration, name, options)
+          @registry[name] = RegisteredIntegration.new(name, integration, options)
+        end
+
+        def self.registry
+          @registry
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "patcher"
+require_relative "request_middleware"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # Rack integration for AI Guard
+        class Integration
+          include Datadog::AIGuard::Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new("1.1.0")
+
+          register_as :rack, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs["rack"]&.version
+          end
+
+          def self.loaded?
+            !defined?(::Rack).nil?
+          end
+
+          def self.compatible?
+            super && !!(version&.>= MINIMUM_VERSION)
+          end
+
+          def self.auto_instrument?
+            false
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # Patcher for Rack integration
+        module Patcher
+          module_function
+
+          def patched?
+            !!Patcher.instance_variable_get(:@patched)
+          end
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            Patcher.instance_variable_set(:@patched, true)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/contrib/rack/request_middleware.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require_relative "../../../tracing/client_ip"
+require_relative "../../../tracing/contrib/rack/header_collection"
+require_relative "../../../tracing/metadata/ext"
+require_relative "../../ext"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # AI Guard Rack middleware.
+        #
+        # At request entry the middleware stores some request attributes for
+        # on the active trace under the `_dd.ai_guard.` prefix.
+        #
+        # Later when AI Guard evaluation is performed, those attributes are
+        # mirrored on AI Guard span.
+        class RequestMiddleware
+          def initialize(app, opt = {})
+            @app = app
+          end
+
+          def call(env)
+            trace = Datadog::Tracing.active_trace
+            return @app.call(env) unless trace
+
+            store_anomaly_detection_tags!(trace, env)
+
+            @app.call(env)
+          ensure
+            # @type var trace: Datadog::Tracing::TraceSegment?
+            # Steep: https://github.com/soutaro/steep/issues/919
+            if trace
+              tag_client_ip_on_request_span!(trace) if ai_guard_executed?(trace)
+
+              clean_up_ai_guard_temp_tags!(trace)
+            end
+          end
+
+          private
+
+          # steep:ignore:start
+          def store_anomaly_detection_tags!(trace, env)
+            remote_ip = env["REMOTE_ADDR"]
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_NETWORK_CLIENT_IP_TAG, remote_ip) if remote_ip
+
+            headers = Datadog::Tracing::Contrib::Rack::Header::RequestHeaderCollection.new(env)
+            resolved_client_ip = Datadog::Tracing::ClientIp.extract_client_ip(headers, remote_ip)
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_HTTP_CLIENT_IP_TAG, resolved_client_ip) if resolved_client_ip
+
+            user_agent = env["HTTP_USER_AGENT"]
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_HTTP_USERAGENT_TAG, user_agent) if user_agent
+          rescue => e
+            Datadog::AIGuard.telemetry&.report(e, description: "AI Guard: failed to get request attributes")
+          end
+          # steep:ignore:end
+
+          # steep:ignore:start
+          def clean_up_ai_guard_temp_tags!(trace)
+            Ext::TRACE_ANOMALY_DETECTION_TAGS.each do |tag|
+              trace.clear_tag(tag)
+            end
+
+            trace.clear_tag(Datadog::AIGuard::Ext::TRACE_EXECUTED_TAG)
+          end
+          # steep:ignore:end
+
+          # AI Guard's evaluation flow sets `_dd.ai_guard.executed` on the
+          # trace whenever an AI Guard span is created during the request.
+          # steep:ignore:start
+          def ai_guard_executed?(trace)
+            trace.get_tag(Datadog::AIGuard::Ext::TRACE_EXECUTED_TAG) == "1"
+          end
+          # steep:ignore:end
+
+          # steep:ignore:start
+          def tag_client_ip_on_request_span!(trace)
+            span = Datadog::Tracing.active_span
+            return unless span
+
+            client_ip = trace.get_tag(Datadog::AIGuard::Ext::TRACE_HTTP_CLIENT_IP_TAG)
+            if client_ip && span.get_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP).nil?
+              span.set_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP, client_ip)
+            end
+
+            network_client_ip = trace.get_tag(Datadog::AIGuard::Ext::TRACE_NETWORK_CLIENT_IP_TAG)
+            span["network.client.ip"] = network_client_ip if network_client_ip
+          rescue => e
+            Datadog::AIGuard.telemetry&.report(e, description: "AI Guard: failed to tag client IP on root span")
+          end
+          # steep:ignore:end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "patcher"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rails
+        # Rails integration for AI Guard
+        class Integration
+          include Datadog::AIGuard::Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new("4")
+
+          register_as :rails, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs["railties"]&.version
+          end
+
+          def self.loaded?
+            !defined?(::Rails).nil?
+          end
+
+          def self.compatible?
+            super && !!(version&.>= MINIMUM_VERSION)
+          end
+
+          def self.auto_instrument?
+            true
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require_relative "../../../core/utils/only_once"
+require_relative "../rack/request_middleware"
+require_relative "../../../tracing/contrib"
+require_relative "../../../tracing/contrib/rack/middlewares"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rails
+        # Patcher for AI Guard on Rails. Inserts the AI Guard Rack middleware
+        # right after the Tracing Rack middleware so the request span is
+        # already active when AI Guard tags the client IP.
+        module Patcher
+          BEFORE_INITIALIZE_ONLY_ONCE_PER_APP = Hash.new { |h, key| h[key] = Datadog::Core::Utils::OnlyOnce.new }
+
+          module_function
+
+          def patched?
+            !!Patcher.instance_variable_get(:@patched)
+          end
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            patch_before_initialize
+            Patcher.instance_variable_set(:@patched, true)
+          end
+
+          def patch_before_initialize
+            ::ActiveSupport.on_load(:before_initialize) do
+              Datadog::AIGuard::Contrib::Rails::Patcher.before_initialize(self)
+            end
+          end
+
+          def before_initialize(app)
+            BEFORE_INITIALIZE_ONLY_ONCE_PER_APP[app].run do
+              # Middleware must be added before the application is initialized.
+              # Otherwise the middleware stack will be frozen.
+              add_middleware(app) if Datadog.configuration.tracing[:rails][:middleware]
+            end
+          end
+
+          def add_middleware(app)
+            if include_middleware?(Datadog::Tracing::Contrib::Rack::TraceMiddleware, app)
+              app.middleware.insert_after(
+                Datadog::Tracing::Contrib::Rack::TraceMiddleware,
+                Datadog::AIGuard::Contrib::Rack::RequestMiddleware
+              )
+            else
+              app.middleware.insert_before(0, Datadog::AIGuard::Contrib::Rack::RequestMiddleware)
+            end
+          end
+
+          def include_middleware?(middleware, app)
+            found = false
+
+            # find tracer middleware reference in Rails::Configuration::MiddlewareStackProxy
+            app.middleware.instance_variable_get(:@operations).each do |operation|
+              args = case operation
+              when Array
+                # rails 5.2
+                _op, args = operation
+                args
+              when Proc
+                if operation.binding.local_variables.include?(:args)
+                  # rails 6.0, 6.1
+                  operation.binding.local_variable_get(:args)
+                else
+                  # rails 7.0 uses ... to pass args
+                  # steep:ignore:start
+                  args_getter = Class.new do
+                    def method_missing(_op, *args) # standard:disable Style/MissingRespondToMissing
+                      args
+                    end
+                  end.new
+                  # steep:ignore:end
+                  operation.call(args_getter)
+                end
+              else
+                # unknown, pass through
+                []
+              end
+
+              found = true if args.include?(middleware)
+            end
+
+            found
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module RubyLLM
+        # module that gets prepended to RubyLLM::Chat
+        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
+
+              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
+
+          def complete(&block)
+            Datadog::AIGuard::Contrib::RubyLLM::ChatInstrumentation.evaluate!(messages)
+
+            super
+          end
+
+          def handle_tool_calls(response, &block)
+            Datadog::AIGuard::Contrib::RubyLLM::ChatInstrumentation.evaluate!(messages)
+
+            super
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "patcher"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module RubyLLM
+        # This class provides helper methods that are used when patching RubyLLM
+        class Integration
+          include Datadog::AIGuard::Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new("1.0.0")
+
+          register_as :ruby_llm, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs["ruby_llm"]&.version
+          end
+
+          def self.loaded?
+            !defined?(::RubyLLM).nil?
+          end
+
+          def self.compatible?
+            super && !!(version&.>= MINIMUM_VERSION)
+          end
+
+          def self.auto_instrument?
+            false
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require_relative "chat_instrumentation"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module RubyLLM
+        # AIGuard patcher module for RubyLLM
+        module Patcher
+          module_function
+
+          def patched?
+            !!Patcher.instance_variable_get(:@patched)
+          end
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            ::RubyLLM::Chat.prepend(ChatInstrumentation)
+
+            Patcher.instance_variable_set(:@patched, true)
+          end
+        end
+      end
+    end
+  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/message.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Message class for AI Guard
+      class Message
+        attr_reader :role, :content, :tool_call, :tool_call_id
+
+        def initialize(role:, content: nil, tool_call: nil, tool_call_id: nil)
+          raise ArgumentError, "Role must be set to a non-empty value" if role.to_s.empty?
+
+          @role = role.to_sym
+          @content = content
+          @tool_call = tool_call
+          @tool_call_id = tool_call_id
+
+          if @tool_call && !@tool_call.is_a?(ToolCall)
+            raise ArgumentError, "Expected an instance of #{ToolCall.name} for :tool_call argument"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    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
+
+        def initialize
+          @action = Result::ALLOW_ACTION
+          @reason = "AI Guard is disabled"
+          @tags = []
+          @sds_findings = []
+          @tag_probabilities = {}
+        end
+
+        def allow?
+          true
+        end
+
+        def deny?
+          false
+        end
+
+        def abort?
+          false
+        end
+
+        def blocking_enabled?
+          false
+        end
+      end
+    end
+  end
+end