datadog 2.17.0 → 2.37.0

data/lib/datadog/ai_guard/configuration.rb

@@ -0,0 +1,115 @@
+# 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
+                o.skip_telemetry true
+              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

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

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      # Auto-instrumentation for AI Guard integrations
+      module AutoInstrument
+        def self.patch_all
+          integrations = []
+
+          Datadog::AIGuard::Contrib::Integration.registry.each_value do |integration|
+            next unless integration.klass.auto_instrument?
+
+            integrations << integration.name
+          end
+
+          integrations.each do |integration_name|
+            Datadog.configuration.ai_guard.instrument(integration_name)
+          end
+        end
+      end
+    end
+  end
+end

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