OpenAI: instrument embeddings and chat completions

.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

lib/new_relic/agent/configuration/default_source.rb

@@ -1570,6 +1570,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/instrumentation/net_http/instrumentation.rb

@@ -7,6 +7,7 @@ module Agent
     module Instrumentation
       module NetHTTP
         INSTRUMENTATION_NAME = NewRelic::Agent.base_name(name)
+        OPENAI_SEGMENT_PATTERN = %r{Llm/.*/OpenAI/.*}.freeze
 
         def request_with_tracing(request)
           NewRelic::Agent.record_instrumentation_invocation(INSTRUMENTATION_NAME)
@@ -32,12 +33,28 @@ def request_with_tracing(request)
             end
 
             wrapped_response = NewRelic::Agent::HTTPClients::NetHTTPResponse.new(response)
+
+            if openai_parent?(segment)
+              populate_openai_response_headers(wrapped_response, segment.parent)
+            end
+
             segment.process_response_headers(wrapped_response)
+
             response
           ensure
             segment&.finish
           end
         end
+
+        def openai_parent?(segment)
+          segment&.parent&.name&.match?(OPENAI_SEGMENT_PATTERN)
+        end
+
+        def populate_openai_response_headers(response, parent)
+          return unless parent.instance_variable_defined?(:@llm_event)
+
+          parent.llm_event.populate_openai_response_headers(response.to_hash)
+        end
       end
     end
   end

lib/new_relic/agent/instrumentation/ruby_openai.rb

@@ -0,0 +1,37 @@
+# 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'
+
+  OPENAI_VERSION = Gem::Version.new(OpenAI::VERSION) if defined?(OpenAI)
+
+  depends_on do
+    # add a config check for ai_monitoring.enabled
+    # maybe add DT check here eventually?
+    defined?(OpenAI) && defined?(OpenAI::Client) &&
+      OPENAI_VERSION >= Gem::Version.new('3.4.0')
+  end
+
+  executes do
+    if use_prepend?
+      if 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

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

@@ -0,0 +1,181 @@
+# 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
+    VENDOR = 'openAI' # AIM expects this capitalization style for the UI
+    INSTRUMENTATION_NAME = NewRelic::Agent.base_name(name)
+    EMBEDDINGS_PATH = '/embeddings'
+    CHAT_COMPLETIONS_PATH = '/chat/completions'
+    EMBEDDINGS_SEGMENT_NAME = 'Llm/embedding/OpenAI/embeddings'
+    CHAT_COMPLETIONS_SEGMENT_NAME = 'Llm/completion/OpenAI/chat'
+
+    def json_post_with_new_relic(path:, parameters:)
+      return yield unless path == EMBEDDINGS_PATH || path == CHAT_COMPLETIONS_PATH
+
+      NewRelic::Agent.record_instrumentation_invocation(INSTRUMENTATION_NAME)
+      NewRelic::Agent::Llm::LlmEvent.set_llm_agent_attribute_on_transaction
+
+      if path == EMBEDDINGS_PATH
+        embeddings_instrumentation(parameters) { yield }
+      elsif path == CHAT_COMPLETIONS_PATH
+        chat_completions_instrumentation(parameters) { yield }
+      end
+    end
+
+    private
+
+    def embeddings_instrumentation(parameters)
+      segment = NewRelic::Agent::Tracer.start_segment(name: EMBEDDINGS_SEGMENT_NAME)
+      record_openai_metric
+      event = create_embeddings_event(parameters)
+      segment.llm_event = event
+      begin
+        response = NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+        # TODO: Remove !response.include?('error) when we drop support for versions below 4.0.0
+        add_embeddings_response_params(response, event) if response && !response.include?('error')
+
+        response
+      ensure
+        finish(segment, event)
+      end
+    end
+
+    def chat_completions_instrumentation(parameters)
+      segment = NewRelic::Agent::Tracer.start_segment(name: CHAT_COMPLETIONS_SEGMENT_NAME)
+      record_openai_metric
+      event = create_chat_completion_summary(parameters)
+      segment.llm_event = event
+      messages = create_chat_completion_messages(parameters, event.id)
+
+      begin
+        response = NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+        # TODO: Remove !response.include?('error) when we drop support for versions below 4.0.0
+        if response && !response.include?('error')
+          add_chat_completion_response_params(parameters, response, event)
+          messages = update_chat_completion_messages(messages, response, event)
+        end
+
+        response
+      ensure
+        finish(segment, event)
+        messages&.each { |m| m.record }
+      end
+    end
+
+    def create_chat_completion_summary(parameters)
+      NewRelic::Agent::Llm::ChatCompletionSummary.new(
+        # TODO: POST-GA: Add metadata from add_custom_attributes if prefixed with 'llm.', except conversation_id
+        vendor: VENDOR,
+        conversation_id: conversation_id,
+        api_key_last_four_digits: parse_api_key,
+        request_max_tokens: parameters[:max_tokens] || parameters['max_tokens'],
+        request_model: parameters[:model] || parameters['model'],
+        temperature: parameters[:temperature] || parameters['temperature']
+      )
+    end
+
+    def create_embeddings_event(parameters)
+      NewRelic::Agent::Llm::Embedding.new(
+        # TODO: POST-GA: Add metadata from add_custom_attributes if prefixed with 'llm.', except conversation_id
+        vendor: VENDOR,
+        input: parameters[:input] || parameters['input'],
+        api_key_last_four_digits: parse_api_key,
+        request_model: parameters[:model] || parameters['model']
+      )
+    end
+
+    def add_chat_completion_response_params(parameters, response, event)
+      event.response_number_of_messages = (parameters[:messages] || parameters['messages']).size + response['choices'].size
+      # The response hash always returns keys as strings, so we don't need to run an || check here
+      event.response_model = response['model']
+      event.response_usage_total_tokens = response['usage']['total_tokens']
+      event.response_usage_prompt_tokens = response['usage']['prompt_tokens']
+      event.response_usage_completion_tokens = response['usage']['completion_tokens']
+      event.response_choices_finish_reason = response['choices'][0]['finish_reason']
+    end
+
+    def add_embeddings_response_params(response, event)
+      event.response_model = response['model']
+      event.response_usage_total_tokens = response['usage']['total_tokens']
+      event.response_usage_prompt_tokens = response['usage']['prompt_tokens']
+    end
+
+    def parse_api_key
+      'sk-' + headers['Authorization'][-4..-1]
+    end
+
+    # The customer must call add_custom_attributes with llm.conversation_id
+    # before the transaction starts. Otherwise, the conversation_id will be nil.
+    def conversation_id
+      return @nr_conversation_id if @nr_conversation_id
+
+      @nr_conversation_id ||= NewRelic::Agent::Tracer.current_transaction.attributes.custom_attributes[NewRelic::Agent::Llm::LlmEvent::CUSTOM_ATTRIBUTE_CONVERSATION_ID]
+    end
+
+    def create_chat_completion_messages(parameters, summary_id)
+      (parameters[:messages] || parameters['messages']).map.with_index do |message, index|
+        NewRelic::Agent::Llm::ChatCompletionMessage.new(
+          content: message[:content] || message['content'],
+          role: message[:role] || message['role'],
+          sequence: index,
+          completion_id: summary_id,
+          vendor: VENDOR,
+          is_response: true
+        )
+      end
+    end
+
+    def create_chat_completion_response_messages(response, sequence_origin, summary_id)
+      response['choices'].map.with_index(sequence_origin) do |choice, index|
+        NewRelic::Agent::Llm::ChatCompletionMessage.new(
+          content: choice['message']['content'],
+          role: choice['message']['role'],
+          sequence: index,
+          completion_id: summary_id,
+          vendor: VENDOR,
+          is_response: true
+        )
+      end
+    end
+
+    def update_chat_completion_messages(messages, response, summary)
+      messages += create_chat_completion_response_messages(response, messages.size, summary.id)
+      response_id = response['id'] || NewRelic::Agent::GuidGenerator.generate_guid
+
+      messages.each do |message|
+        # TODO: POST-GA: Add metadata from add_custom_attributes if prefixed with 'llm.', except conversation_id
+        message.id = "#{response_id}-#{message.sequence}"
+        message.conversation_id = conversation_id
+        message.request_id = summary.request_id
+        message.response_model = response['model']
+      end
+    end
+
+    def record_openai_metric
+      NewRelic::Agent.record_metric(nr_supportability_metric, 0.0)
+    end
+
+    def segment_noticed_error?(segment)
+      segment&.instance_variable_get(:@noticed_error)
+    end
+
+    def nr_supportability_metric
+      @nr_supportability_metric ||= "Supportability/Ruby/ML/OpenAI/#{::OpenAI::VERSION}"
+    end
+
+    def finish(segment, event)
+      segment&.finish
+
+      return unless event
+
+      if segment
+        event.error = true if segment_noticed_error?(segment)
+        event.duration = segment.duration
+      end
+
+      event.record
+    end
+  end
+end

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

@@ -0,0 +1,20 @@
+# 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::Prepend
+    include NewRelic::Agent::Instrumentation::OpenAI
+
+    # In versions 4.0.0+ json_post is an instance method defined in the
+    # OpenAI::HTTP module, included by the OpenAI::Client class.
+    #
+    # In versions below 4.0.0 json_post is a class method on OpenAI::Client.
+    #
+    # Dependency detection will apply the instrumentation to the correct scope,
+    # so we don't need to change the code here.
+    def json_post(**kwargs)
+      json_post_with_new_relic(**kwargs) { super }
+    end
+  end
+end

lib/new_relic/agent/llm/chat_completion_summary.rb

@@ -12,10 +12,18 @@ class ChatCompletionSummary < LlmEvent
         include ResponseHeaders
 
         ATTRIBUTES = %i[api_key_last_four_digits request_max_tokens
-          response_number_of_messages request_model response_organization
-          response_usage_total_tokens response_usage_prompt_tokens
-          response_usage_completion_tokens response_choices_finish_reason
+          response_number_of_messages request_model response_usage_total_tokens response_usage_prompt_tokens response_usage_completion_tokens response_choices_finish_reason
           request_temperature duration error]
+        ATTRIBUTE_NAME_EXCEPTIONS = {
+          response_number_of_messages: 'response.number_of_messages',
+          request_model: 'request.model',
+          response_usage_total_tokens: 'response.usage.total_tokens',
+          response_usage_prompt_tokens: 'response.usage.prompt_tokens',
+          response_usage_completion_tokens: 'response.usage.completion_tokens',
+          response_choices_finish_reason: 'response.choices.finish_reason',
+          temperature: 'request.temperature'
+        }
+
         EVENT_NAME = 'LlmChatCompletionSummary'
 
         attr_accessor(*ATTRIBUTES)
@@ -24,6 +32,17 @@ def attributes
           LlmEvent::ATTRIBUTES + ChatCompletion::ATTRIBUTES + ResponseHeaders::ATTRIBUTES + ATTRIBUTES
         end
 
+        def attribute_name_exceptions
+          # TODO: OLD RUBIES < 2.6
+          # Hash#merge accepts multiple arguments in 2.6
+          # Remove condition once support for Ruby <2.6 is dropped
+          if RUBY_VERSION >= '2.6.0'
+            LlmEvent::ATTRIBUTE_NAME_EXCEPTIONS.merge(ResponseHeaders::ATTRIBUTE_NAME_EXCEPTIONS, ATTRIBUTE_NAME_EXCEPTIONS)
+          else
+            LlmEvent::ATTRIBUTE_NAME_EXCEPTIONS.merge(ResponseHeaders::ATTRIBUTE_NAME_EXCEPTIONS).merge(ATTRIBUTE_NAME_EXCEPTIONS)
+          end
+        end
+
         def event_name
           EVENT_NAME
         end