OpenAI: instrument embeddings and chat completions

.github/workflows/ci.yml

@@ -185,7 +185,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

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

@@ -32,12 +32,31 @@ def request_with_tracing(request)
             end
 
             wrapped_response = NewRelic::Agent::HTTPClients::NetHTTPResponse.new(response)
+            add_llm_response_headers(wrapped_response, segment.parent) if llm_parent?(segment)
             segment.process_response_headers(wrapped_response)
             response
           ensure
             segment&.finish
           end
         end
+
+        def llm_parent?(segment)
+          segment&.parent&.name&.match?(/Llm\/.*\/OpenAI\/create/)
+        end
+
+        def add_llm_response_headers(response, parent)
+          return unless parent.instance_variable_defined?(:@chat_completion_summary) || parent.instance_variable_defined?(:@embedding) # and maybe log a warning??
+
+          if parent.instance_variable_defined?(:@chat_completion_summary)
+            event = parent.chat_completion_summary
+            event.request_id = response[NewRelic::Agent::Llm::LlmEvent::X_REQUEST_ID] # every event needs this, maybe we should move it someplace else?
+            event.populate_openai_response_headers(response.to_hash)
+          elsif parent.instance_variable_defined?(:@embedding)
+            event = parent.embedding
+            event.request_id = response[NewRelic::Agent::Llm::LlmEvent::X_REQUEST_ID] # every event needs this, maybe we should move it someplace else?
+            event.populate_openai_response_headers(response.to_hash)
+          end
+        end
       end
     end
   end

lib/new_relic/agent/instrumentation/ruby_openai.rb

@@ -0,0 +1,33 @@
+# 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
+    defined?(OpenAI) && defined?(OpenAI::Client)
+    # maybe add DT check here eventually?
+    # possibly also a config check for ai.enabled
+  end
+
+  executes do
+    NewRelic::Agent.logger.info('Installing ruby-openai instrumentation')
+
+    if use_prepend?
+      # instead of metaprogramming on OpenAI::Client, we could also use OpenAI::HTTP,
+      # it's a module that's required by OpenAI::Client and contains the
+      # json_post method we're instrumenting
+      prepend_instrument OpenAI::Client,
+        NewRelic::Agent::Instrumentation::OpenAI::Prepend,
+        NewRelic::Agent::Instrumentation::OpenAI::VENDOR
+    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,21 @@
+# 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)
+
+        def json_post(**kwargs)
+          json_post_with_new_relic(**kwargs) do
+            json_post_without_new_relic(**kwargs)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,180 @@
+# 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' # or SUPPORTBILITY_NAME? or both?
+    # TODO: should everything below be called embeddings if we renamed to chat completions?
+    EMBEDDINGS_PATH = '/embeddings'
+    CHAT_COMPLETIONS_PATH = '/chat/completions'
+    SEGMENT_NAME_FORMAT = 'Llm/%s/OpenAI/create'
+
+    # This method is defined in the OpenAI::HTTP module that is included
+    # only in the OpenAI::Client class
+    def json_post_with_new_relic(path:, parameters:)
+      if path == EMBEDDINGS_PATH
+        NewRelic::Agent.record_instrumentation_invocation(VENDOR)
+        embedding_instrumentation(parameters) { yield }
+      elsif path == CHAT_COMPLETIONS_PATH
+        NewRelic::Agent.record_instrumentation_invocation(VENDOR)
+        chat_completions_instrumentation(parameters) { yield }
+      else
+        yield
+      end
+    end
+
+    private
+
+    def embedding_instrumentation(parameters)
+      segment = NewRelic::Agent::Tracer.start_segment(SEGMENT_NAME_FORMAT % 'embedding')
+      record_openai_metric
+      event = create_embedding_event(parameters)
+      segment.embedding = event
+      begin
+        response = NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+
+        response
+      ensure
+        add_embedding_response_params(response, event) if response
+        segment&.finish
+        event&.error = true if segment_noticed_error?(segment) # need to test throwing an error
+        event&.duration = segment&.duration
+        event&.record # always record the event
+      end
+    end
+
+    def chat_completions_instrumentation(parameters)
+      # TODO: Do we have to start the segment outside the ensure block?
+      segment = NewRelic::Agent::Tracer.start_segment(name: SEGMENT_NAME_FORMAT % 'completion')
+      record_openai_metric
+      event = create_chat_completion_summary(parameters)
+      segment.chat_completion_summary = event
+      messages = create_chat_completion_messages(parameters, summary_event_id)
+      response = NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+      add_response_params(parameters, response, event) if response
+      messages = update_chat_completion_messages(messages, response, event) if response
+
+      response # return the response to the original caller
+    ensure
+      segment&.finish
+      event&.error = true if segment_noticed_error?(segment)
+      event&.duration = segment&.duration
+      event&.record # always record the event
+      messages&.each { |m| m&.record }
+    end
+
+    def create_chat_completion_summary(parameters)
+      event = NewRelic::Agent::Llm::ChatCompletionSummary.new(
+        # metadata => TBD, create API
+        vendor: VENDOR,
+        conversation_id: conversation_id,
+        api_key_last_four_digits: parse_api_key,
+        # TODO: Determine how to access parameters with keys as strings
+        request_max_tokens: parameters[:max_tokens],
+        request_model: parameters[:model],
+        temperature: parameters[:temperature]
+      )
+    end
+
+    def create_embedding_event(parameters)
+      # TODO: Determine how to access parameters with keys as strings
+      event = NewRelic::Agent::Llm::Embedding.new(
+        # metadata => TBD, create API
+        vendor: VENDOR,
+        input: parameters[:input],
+        api_key_last_four_digits: parse_api_key,
+        request_model: parameters[:model]
+      )
+    end
+
+    def add_response_params(parameters, response, event)
+      event.response_number_of_messages = parameters[:messages].size + response['choices'].size
+      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_embedding_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 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['conversation_id']
+    end
+
+    def create_chat_completion_messages(parameters, summary_id)
+      # TODO: Determine how to access parameters with keys as strings
+      parameters[:messages].map.with_index do |message, i|
+        NewRelic::Agent::Llm::ChatCompletionMessage.new(
+          content: message[:content] || message['content'],
+          role: message[:role] || message['role'],
+          sequence: i,
+          completion_id: summary_id,
+          vendor: VENDOR,
+          is_response: false
+        )
+      end
+    end
+
+    def create_chat_completion_response_messages(response, sequence_origin, summary_id)
+      response['choices'].map.with_index(sequence_origin) do |choice, i|
+        NewRelic::Agent::Llm::ChatCompletionMessage.new(
+          content: choice['message']['content'],
+          role: choice['message']['role'],
+          sequence: i,
+          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|
+        # metadata => TBD, create API
+        message.id = "#{response_id}-#{message.sequence}"
+        message.conversation_id = conversation_id
+        message.request_id = summary.request_id
+        message.response_model = response['model']
+      end
+    end
+
+    # Name is defined in Ruby 3.0+
+    # copied from rails code
+    # Parameter keys might be symbols and might be strings
+    # response body keys have always been strings
+    def hash_with_indifferent_access_whatever
+      if Symbol.method_defined?(:name)
+        key.kind_of?(Symbol) ? key.name : key
+      else
+        key.kind_of?(Symbol) ? key.to_s : key
+      end
+    end
+
+    # the preceding :: are necessary to access the OpenAI module defined in the gem rather than the current module
+    # TODO: discover whether this metric name should be prepended with 'Supportability'
+    def record_openai_metric
+      NewRelic::Agent.record_metric("Ruby/ML/OpenAI/#{::OpenAI::VERSION}", 0.0)
+    end
+
+    def segment_noticed_error?(segment)
+      segment&.instance_variable_get(:@noticed_error)
+    end
+  end
+end

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

@@ -0,0 +1,13 @@
+# 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
+
+    def json_post(**kwargs)
+      json_post_with_new_relic(**kwargs) { super }
+    end
+  end
+end

lib/new_relic/agent/llm/chat_completion_summary.rb

@@ -16,6 +16,16 @@ class ChatCompletionSummary < LlmEvent
           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 +34,10 @@ def attributes
           LlmEvent::ATTRIBUTES + ChatCompletion::ATTRIBUTES + ResponseHeaders::ATTRIBUTES + ATTRIBUTES
         end
 
+        def attribute_name_exceptions
+          LlmEvent::ATTRIBUTE_NAME_EXCEPTIONS.merge(ResponseHeaders::ATTRIBUTE_NAME_EXCEPTIONS, ATTRIBUTE_NAME_EXCEPTIONS)
+        end
+
         def event_name
           EVENT_NAME
         end

lib/new_relic/agent/llm/embedding.rb

@@ -9,8 +9,13 @@ class Embedding < LlmEvent
         include ResponseHeaders
 
         ATTRIBUTES = %i[input api_key_last_four_digits request_model
-          response_organization response_usage_total_tokens
-          response_usage_prompt_tokens duration error]
+          response_usage_total_tokens response_usage_prompt_tokens duration
+          error]
+        ATTRIBUTE_NAME_EXCEPTIONS = {
+          request_model: 'request.model',
+          response_usage_total_tokens: 'response.usage.total_tokens',
+          response_usage_prompt_tokens: 'response.usage.prompt_tokens'
+        }
         EVENT_NAME = 'LlmEmbedding'
 
         attr_accessor(*ATTRIBUTES)
@@ -19,6 +24,10 @@ def attributes
           LlmEvent::ATTRIBUTES + ResponseHeaders::ATTRIBUTES + ATTRIBUTES
         end
 
+        def attribute_name_exceptions
+          LlmEvent::ATTRIBUTE_NAME_EXCEPTIONS.merge(ResponseHeaders::ATTRIBUTE_NAME_EXCEPTIONS, ATTRIBUTE_NAME_EXCEPTIONS)
+        end
+
         def event_name
           EVENT_NAME
         end