Merge pull request #2442 from newrelic/openai_instrumentation

AIM Instrumentation - OpenAI

.github/workflows/ci.yml

@@ -189,7 +189,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        multiverse: [agent, background, background_2, database, frameworks, httpclients, httpclients_2, rails, rest]
+        multiverse: [agent, ai, background, background_2, database, frameworks, httpclients, httpclients_2, rails, rest]
         ruby-version: [2.4.10, 3.3.0]
 
     steps:

.github/workflows/ci_cron.yml

@@ -199,7 +199,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        multiverse: [agent, background, background_2, database, frameworks, httpclients, httpclients_2, rails, rest]
+        multiverse: [agent, ai, background, background_2, database, frameworks, httpclients, httpclients_2, rails, rest]
         ruby-version: [2.4.10, 2.5.9, 2.6.10, 2.7.8, 3.0.6, 3.1.4, 3.2.2, 3.3.0]
     steps:
       - name: Configure git

.github/workflows/ci_jruby.yml

@@ -39,7 +39,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        suite: [active_support_broadcast_logger, active_support_logger, activemerchant, agent_only, async_http, bare, deferred_instrumentation, grape, high_security, httpclient, httprb, httpx, json, logger, marshalling, rack, resque, roda, roda_agent_disabled, sequel, sinatra, sinatra_agent_disabled, stripe, thread, tilt, typhoeus]
+        suite: [active_support_broadcast_logger, active_support_logger, activemerchant, agent_only, async_http, bare, deferred_instrumentation, grape, high_security, httpclient, httprb, httpx, json, logger, marshalling, rack, resque, roda, roda_agent_disabled, ruby_openai, sequel, sinatra, sinatra_agent_disabled, stripe, thread, tilt, typhoeus]
 
     steps:
       - name: Configure git

CHANGELOG.md

@@ -4,6 +4,20 @@
 
 Version <dev> introduces the option to store tracer state on the thread-level and hardens the browser agent insertion logic to better proactively anticipate errors.
 
+- **Feature: Add instrumentation for ruby-openai**
+
+  Instrumentation has been added for the [ruby-openai](https://github.com/alexrudall/ruby-openai) gem, supporting versions 3.4.0 and higher [(PR#2442)](https://github.com/newrelic/newrelic-ruby-agent/pull/2442). While ruby-openai instrumentation is enabled by default, the configuration option `ai_monitoring.enabled` is disabled by default and controls all AI monitoring. `ai_monitoring.enabled` must be set to `true` in order to receive ruby-openai instrumentation. High-Security Mode must be disabled in order to receive AI monitoring.
+
+  Calls to embedding and chat completion endpoints are automatically traced. These events can be enhanced with the introduction of two new APIs. Custom attributes can also be added to LLM events using the API `NewRelic::Agent.add_custom_attributes`, but they must be prefixed with `llm.`. For example, `NewRelic::Agent.add_custom_attributes({'llm.user_id': user_id})`.
+
+- **Feature: Add AI monitoring APIs**
+
+  This version introduces two new APIs that allow users to record additional information on LLM events:
+  * `NewRelic::Agent.record_llm_feedback_event` - Records user feedback events.
+  * `NewRelic::Agent.set_llm_token_count_callback` - Sets a callback proc for calculating `token_count` attributes for embedding and chat completion message events.
+
+  Visit [RubyDoc](https://rubydoc.info/github/newrelic/newrelic-ruby-agent/) for more information on each of these APIs.
+
 - **Feature: Store tracer state on thread-level**
 
   A new configuration option, `thread_local_tracer_state`, stores New Relic's tracer state on the thread-level, as opposed to the default fiber-level storage. This configuration is turned off by default. Our thanks go to community member [@markiz](https://github.com/markiz) who contributed the idea, code, configuration option, and tests for this new feature! [PR#2475](https://github.com/newrelic/newrelic-ruby-agent/pull/2475).

lib/new_relic/agent.rb

@@ -107,11 +107,14 @@ class SerializationError < StandardError; end
 
     # placeholder name used when we cannot determine a transaction's name
     UNKNOWN_METRIC = '(unknown)'.freeze
+    LLM_FEEDBACK_MESSAGE = 'LlmFeedbackMessage'
 
     attr_reader :error_group_callback
+    attr_reader :llm_token_count_callback
 
     @agent = nil
     @error_group_callback = nil
+    @llm_token_count_callback = nil
     @logger = nil
     @tracer_lock = Mutex.new
     @tracer_queue = []
@@ -389,6 +392,92 @@ def record_custom_event(event_type, event_attrs)
       nil
     end
 
+    # Records user feedback events for LLM applications. This API must pass
+    # the current trace id as a parameter, which can be obtained using:
+    #
+    #   NewRelic::Agent::Tracer.current_trace_id
+    #
+    # @param [String] ID of the trace where the chat completion(s) related
+    #   to the feedback occurred.
+    #
+    # @param [String or Integer] Rating provided by an end user
+    #   (ex: “Good", "Bad”, 1, 2, 5, 8, 10).
+    #
+    # @param [optional, String] Category of the feedback as provided by the
+    #   end user (ex: “informative”, “inaccurate”).
+    #
+    # @param start_time [optional, String] Freeform text feedback from an
+    #   end user.
+    #
+    # @param [optional, Hash] Set of key-value pairs to store any other
+    #   desired data to submit with the feedback event.
+    #
+    # @api public
+    #
+    def record_llm_feedback_event(trace_id:,
+      rating:,
+      category: nil,
+      message: nil,
+      metadata: NewRelic::EMPTY_HASH)
+
+      record_api_supportability_metric(:record_llm_feedback_event)
+      unless NewRelic::Agent.config[:'distributed_tracing.enabled']
+        return NewRelic::Agent.logger.error('Distributed tracing must be enabled to record LLM feedback')
+      end
+
+      feedback_message_event = {
+        'trace_id': trace_id,
+        'rating': rating,
+        'category': category,
+        'message': message,
+        'id': NewRelic::Agent::GuidGenerator.generate_guid,
+        'ingest_source': NewRelic::Agent::Llm::LlmEvent::INGEST_SOURCE
+      }
+      feedback_message_event.merge!(metadata) unless metadata.empty?
+
+      NewRelic::Agent.record_custom_event(LLM_FEEDBACK_MESSAGE, feedback_message_event)
+    rescue ArgumentError
+      raise
+    rescue => exception
+      NewRelic::Agent.logger.error('record_llm_feedback_event', exception)
+    end
+
+    # @!endgroup
+
+    # @!group LLM callbacks
+
+    # Set a callback proc for calculating `token_count` attributes for
+    # LlmEmbedding and LlmChatCompletionMessage events
+    #
+    # @param callback_proc [Proc] the callback proc
+    #
+    # This method should be called only once to set a callback for
+    # use with all LLM token calculations. If it is called multiple times, each
+    # new callback will replace the old one.
+    #
+    # The proc will be called with a single hash as its input argument and
+    # must return an Integer representing the number of tokens used for that
+    # particular prompt, completion message, or embedding. Values less than or
+    # equal to 0 will not be attached to an event.
+    #
+    # The hash has the following keys:
+    #
+    # :model => [String] The name of the LLM model
+    # :content => [String] The message content or prompt
+    #
+    # @api public
+    #
+    def set_llm_token_count_callback(callback_proc)
+      unless callback_proc.is_a?(Proc)
+        NewRelic::Agent.logger.error("#{self}.#{__method__}: expected an argument of type Proc, " \
+                                     "got #{callback_proc.class}")
+        return
+      end
+
+      record_api_supportability_metric(:set_llm_token_count_callback)
+      @llm_token_count_callback = callback_proc
+    end
+
     # @!endgroup
 
     # @!group Manual agent configuration and startup/shutdown

lib/new_relic/agent/configuration/default_source.rb

@@ -360,6 +360,26 @@ def self.enforce_fallback(allowed_values: nil, fallback: nil)
               - a.third.event
           DESCRIPTION
         },
+        :'ai_monitoring.enabled' => {
+          :default => false,
+          :public => true,
+          :type => Boolean,
+          :allowed_from_server => false,
+          :description => 'If `false`, all LLM instrumentation (OpenAI only for now) will be disabled and no metrics, events, or spans will be sent. AI Monitoring is automatically disabled if `high_security` mode is enabled.'
+        },
+        :'ai_monitoring.record_content.enabled' => {
+          :default => true,
+          :public => true,
+          :type => Boolean,
+          :allowed_from_server => false,
+          :description => <<~DESCRIPTION
+            If `false`, LLM instrumentation (OpenAI only for now) will not capture input and output content on specific LLM events.
+
+            The excluded attributes include:
+              * `content` from LlmChatCompletionMessage events
+              * `input` from LlmEmbedding events
+          DESCRIPTION
+        },
         # this is only set via server side config
         :apdex_t => {
           :default => 0.5,
@@ -1577,6 +1597,15 @@ def self.enforce_fallback(allowed_values: nil, fallback: nil)
           :allowed_from_server => false,
           :description => 'Controls auto-instrumentation of `Net::HTTP` at start-up. May be one of: `auto`, `prepend`, `chain`, `disabled`.'
         },
+        :'instrumentation.ruby_openai' => {
+          :default => 'auto',
+          :documentation_default => 'auto',
+          :public => true,
+          :type => String,
+          :dynamic_name => true,
+          :allowed_from_server => false,
+          :description => 'Controls auto-instrumentation of the ruby-openai gem at start-up. May be one of: `auto`, `prepend`, `chain`, `disabled`.'
+        },
         :'instrumentation.puma_rack' => {
           :default => value_of(:'instrumentation.rack'),
           :documentation_default => 'auto',

lib/new_relic/agent/configuration/high_security_source.rb

@@ -19,6 +19,7 @@ def initialize(local_settings)
             :'elasticsearch.obfuscate_queries' => true,
             :'transaction_tracer.record_redis_arguments' => false,
 
+            :'ai_monitoring.enabled' => false,
             :'custom_insights_events.enabled' => false,
             :'strip_exception_messages.enabled' => true
           })

lib/new_relic/agent/configuration/security_policy_source.rb

@@ -7,6 +7,8 @@
 module NewRelic
   module Agent
     module Configuration
+      # The Language Security Policy Source gives customers the ability to
+      # configure high security mode settings.
       class SecurityPolicySource < DottedHash
         class << self
           def enabled?(option)
@@ -147,6 +149,15 @@ def change_setting(policies, option, new_value)
               permitted_fn: nil
             }
           ],
+          'ai_monitoring' => [
+            {
+              option: :'ai_monitoring.enabled',
+              supported: true,
+              enabled_fn: method(:enabled?),
+              disabled_value: false,
+              permitted_fn: nil
+            }
+          ],
           'allow_raw_exception_messages' => [
             {
               option: :'strip_exception_messages.enabled',

lib/new_relic/agent/custom_event_aggregator.rb

@@ -52,11 +52,11 @@ def create_event(type, priority, attributes)
           {TYPE => type,
            TIMESTAMP => Process.clock_gettime(Process::CLOCK_REALTIME).to_i,
            PRIORITY => priority},
-          create_custom_event_attributes(attributes)
+          create_custom_event_attributes(type, attributes)
         ]
       end
 
-      def create_custom_event_attributes(attributes)
+      def create_custom_event_attributes(type, attributes)
         result = AttributeProcessing.flatten_and_coerce(attributes)
 
         if result.size > MAX_ATTRIBUTE_COUNT
@@ -70,9 +70,9 @@ def create_custom_event_attributes(attributes)
             key = key[0, MAX_NAME_SIZE]
           end
 
-          # value is limited to 4095
+          # value is limited to 4095 except for LLM content-related events
           if val.is_a?(String) && val.length > MAX_ATTRIBUTE_SIZE
-            val = val[0, MAX_ATTRIBUTE_SIZE]
+            val = val[0, MAX_ATTRIBUTE_SIZE] unless NewRelic::Agent::LLM.exempt_event_attribute?(type, key)
           end
 
           new_result[key] = val

lib/new_relic/agent/error_collector.rb

@@ -216,6 +216,8 @@ def extract_stack_trace(exception)
       def notice_segment_error(segment, exception, options = {})
         return if skip_notice_error?(exception)
 
+        options.merge!(segment.llm_event.error_attributes(exception)) if segment.llm_event
+
         segment.set_noticed_error(create_noticed_error(exception, options))
         exception
       rescue => e

lib/new_relic/agent/instrumentation/net_http/instrumentation.rb

@@ -32,7 +32,13 @@ def request_with_tracing(request)
             end
 
             wrapped_response = NewRelic::Agent::HTTPClients::NetHTTPResponse.new(response)
+
+            if NewRelic::Agent::LLM.openai_parent?(segment)
+              NewRelic::Agent::LLM.populate_openai_response_headers(wrapped_response, segment.parent)
+            end
+
             segment.process_response_headers(wrapped_response)
+
             response
           ensure
             segment&.finish

lib/new_relic/agent/instrumentation/ruby_openai.rb

@@ -0,0 +1,35 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative 'ruby_openai/instrumentation'
+require_relative 'ruby_openai/chain'
+require_relative 'ruby_openai/prepend'
+
+DependencyDetection.defer do
+  named :'ruby_openai'
+
+  depends_on do
+    NewRelic::Agent.config[:'ai_monitoring.enabled'] &&
+      defined?(OpenAI) && defined?(OpenAI::Client) &&
+      Gem::Version.new(OpenAI::VERSION) >= Gem::Version.new('3.4.0')
+  end
+
+  executes do
+    if use_prepend?
+      # TODO: Remove condition when we drop support for versions below 5.0.0
+      if Gem::Version.new(OpenAI::VERSION) >= Gem::Version.new('5.0.0')
+        prepend_instrument OpenAI::Client,
+          NewRelic::Agent::Instrumentation::OpenAI::Prepend,
+          NewRelic::Agent::Instrumentation::OpenAI::VENDOR
+      else
+        prepend_instrument OpenAI::Client.singleton_class,
+          NewRelic::Agent::Instrumentation::OpenAI::Prepend,
+          NewRelic::Agent::Instrumentation::OpenAI::VENDOR
+      end
+    else
+      chain_instrument NewRelic::Agent::Instrumentation::OpenAI::Chain,
+        NewRelic::Agent::Instrumentation::OpenAI::VENDOR
+    end
+  end
+end

lib/new_relic/agent/instrumentation/ruby_openai/chain.rb

@@ -0,0 +1,36 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+module NewRelic::Agent::Instrumentation
+  module OpenAI::Chain
+    def self.instrument!
+      ::OpenAI::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::OpenAI
+
+        alias_method(:json_post_without_new_relic, :json_post)
+
+        # In versions 4.0.0+ json_post is an instance method
+        # defined in the OpenAI::HTTP module, included by the
+        # OpenAI::Client class
+        def json_post(**kwargs)
+          json_post_with_new_relic(**kwargs) do
+            json_post_without_new_relic(**kwargs)
+          end
+        end
+
+        # In versions below 4.0.0 json_post is a class method
+        # on OpenAI::Client
+        class << self
+          alias_method(:json_post_without_new_relic, :json_post)
+
+          def json_post(**kwargs)
+            json_post_with_new_relic(**kwargs) do
+              json_post_without_new_relic(**kwargs)
+            end
+          end
+        end
+      end
+    end
+  end
+end