Merge pull request #2415 from newrelic/llm_events_2

Create LLM event structure (team refactor)

lib/new_relic/agent.rb

@@ -62,6 +62,7 @@ module Agent
     require 'new_relic/agent/attribute_processing'
     require 'new_relic/agent/linking_metadata'
     require 'new_relic/agent/local_log_decorator'
+    require 'new_relic/agent/llm'
 
     require 'new_relic/agent/instrumentation/controller_instrumentation'
 

lib/new_relic/agent/llm.rb

@@ -0,0 +1,11 @@
+# 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 'llm/llm_event'
+require_relative 'llm/chat_completion'
+require_relative 'llm/chat_completion_message'
+require_relative 'llm/chat_completion_summary'
+require_relative 'llm/embedding'
+require_relative 'llm/feedback'
+require_relative 'llm/response_headers'

lib/new_relic/agent/llm/chat_completion.rb

@@ -0,0 +1,15 @@
+# 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
+  module Agent
+    module Llm
+      module ChatCompletion
+        ATTRIBUTES = %i[conversation_id]
+
+        attr_accessor(*ATTRIBUTES)
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/chat_completion_message.rb

@@ -0,0 +1,26 @@
+# 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
+  module Agent
+    module Llm
+      class ChatCompletionMessage < LlmEvent
+        include ChatCompletion
+
+        ATTRIBUTES = %i[content role sequence completion_id is_response]
+        EVENT_NAME = 'LlmChatCompletionMessage'
+
+        attr_accessor(*ATTRIBUTES)
+
+        def attributes
+          LlmEvent::ATTRIBUTES + ChatCompletion::ATTRIBUTES + ATTRIBUTES
+        end
+
+        def event_name
+          EVENT_NAME
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/chat_completion_summary.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 'response_headers'
+
+module NewRelic
+  module Agent
+    module Llm
+      class ChatCompletionSummary < LlmEvent
+        include ChatCompletion
+        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
+          request_temperature duration error]
+        EVENT_NAME = 'LlmChatCompletionSummary'
+
+        attr_accessor(*ATTRIBUTES)
+
+        def attributes
+          LlmEvent::ATTRIBUTES + ChatCompletion::ATTRIBUTES + ResponseHeaders::ATTRIBUTES + ATTRIBUTES
+        end
+
+        def event_name
+          EVENT_NAME
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/embedding.rb

@@ -0,0 +1,28 @@
+# 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
+  module Agent
+    module Llm
+      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]
+        EVENT_NAME = 'LlmEmbedding'
+
+        attr_accessor(*ATTRIBUTES)
+
+        def attributes
+          LlmEvent::ATTRIBUTES + ResponseHeaders::ATTRIBUTES + ATTRIBUTES
+        end
+
+        def event_name
+          EVENT_NAME
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/feedback.rb

@@ -0,0 +1,12 @@
+# 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
+  module Agent
+    module Llm
+      class Feedback < LlmEvent
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/llm_event.rb

@@ -0,0 +1,59 @@
+# 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
+  module Agent
+    module Llm
+      class LlmEvent
+        # Every subclass must define its own ATTRIBUTES constant, an array of symbols representing
+        # that class's unique attributes
+        ATTRIBUTES = %i[id request_id span_id transaction_id
+          trace_id response_model vendor ingest_source]
+        # These attributes should not be passed as arguments to initialize and will be set by the agent
+        AGENT_DEFINED_ATTRIBUTES = %i[span_id transaction_id trace_id ingest_source]
+        INGEST_SOURCE = 'Ruby'
+
+        attr_accessor(*ATTRIBUTES)
+
+        # This initialize method is used for all subclasses.
+        # It leverages the subclass's `attributes` method to iterate through
+        # all the attributes for that subclass.
+        # It assigns instance variables for all arguments passed to the method.
+        # It also assigns agent-defined attributes.
+        def initialize(opts = {})
+          (attributes - AGENT_DEFINED_ATTRIBUTES).each do |attr|
+            instance_variable_set(:"@#{attr}", opts[attr]) if opts.key?(attr)
+          end
+
+          @span_id = NewRelic::Agent::Tracer.current_span_id
+          @transaction_id = NewRelic::Agent::Tracer.current_transaction&.guid
+          @trace_id = NewRelic::Agent::Tracer.current_trace_id
+          @ingest_source = INGEST_SOURCE
+        end
+
+        # All subclasses use event_attributes to get a full hash of all
+        # attributes and their values
+        def event_attributes
+          attributes.each_with_object({}) do |attr, hash|
+            hash[attr] = instance_variable_get(:"@#{attr}")
+          end
+        end
+
+        # Subclasses define an attributes method to concatenate attributes
+        # defined across their ancestors and other modules
+        def attributes
+          ATTRIBUTES
+        end
+
+        # Subclasses that record events will override this method
+        def event_name
+        end
+
+        def record
+          NewRelic::Agent.record_custom_event(event_name, event_attributes)
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/llm/response_headers.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
+  module Agent
+    module Llm
+      module ResponseHeaders
+        ATTRIBUTES = %i[llm_version rate_limit_requests rate_limit_tokens
+          rate_limit_remaining_requests rate_limit_remaining_tokens
+          rate_limit_reset_requests rate_limit_reset_tokens]
+
+        OPENAI_VERSION = 'openai-version'
+        X_RATELIMIT_LIMIT_REQUESTS = 'x-ratelimit-limit-requests'
+        X_RATELIMIT_LIMIT_TOKENS = 'x-ratelimit-limit-tokens'
+        X_RATELIMIT_REMAINING_REQUESTS = 'x-ratelimit-remaining-requests'
+        X_RATELIMIT_REMAINING_TOKENS = 'x-ratelimit-remaining-tokens'
+        X_RATELIMIT_RESET_REQUESTS = 'x-ratelimit-reset-requests'
+        X_RATELIMIT_RESET_TOKENS = 'x-ratelimit-reset-tokens'
+
+        attr_accessor(*ATTRIBUTES)
+
+        # Headers is a hash of Net::HTTP response headers
+        def populate_openai_response_headers(headers)
+          self.llm_version = headers[OPENAI_VERSION]&.first
+          self.rate_limit_requests = headers[X_RATELIMIT_LIMIT_REQUESTS]&.first
+          self.rate_limit_tokens = headers[X_RATELIMIT_LIMIT_TOKENS]&.first
+          self.rate_limit_remaining_requests = headers[X_RATELIMIT_REMAINING_REQUESTS]&.first
+          self.rate_limit_remaining_tokens = headers[X_RATELIMIT_REMAINING_TOKENS]&.first
+          self.rate_limit_reset_requests = headers[X_RATELIMIT_RESET_REQUESTS]&.first
+          self.rate_limit_reset_tokens = headers[X_RATELIMIT_RESET_TOKENS]&.first
+        end
+      end
+    end
+  end
+end

test/new_relic/agent/llm/chat_completion_message_test.rb

@@ -0,0 +1,90 @@
+# 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 '../../../test_helper'
+
+module NewRelic::Agent::Llm
+  class ChatCompletionMessageTest < Minitest::Test
+    def test_attributes_assigned_by_parent_present
+      assert_includes NewRelic::Agent::Llm::ChatCompletionMessage.ancestors, NewRelic::Agent::Llm::LlmEvent
+      assert_includes NewRelic::Agent::Llm::LlmEvent::AGENT_DEFINED_ATTRIBUTES, :transaction_id
+
+      in_transaction do |txn|
+        event = NewRelic::Agent::Llm::ChatCompletionMessage.new
+
+        assert_equal txn.guid, event.transaction_id
+      end
+    end
+
+    def test_attributes_in_parent_list_can_be_assigned_on_init
+      assert_includes NewRelic::Agent::Llm::LlmEvent::ATTRIBUTES, :id
+
+      event = NewRelic::Agent::Llm::ChatCompletionMessage.new(id: 123)
+
+      assert_equal 123, event.id
+    end
+
+    def test_included_module_attributes_list_can_be_assigned_on_init
+      assert_includes NewRelic::Agent::Llm::ChatCompletionMessage.ancestors, NewRelic::Agent::Llm::ChatCompletion
+      assert_includes NewRelic::Agent::Llm::ChatCompletion::ATTRIBUTES, :conversation_id
+
+      conversation_id = '123abc'
+      event = NewRelic::Agent::Llm::ChatCompletionMessage.new(conversation_id: conversation_id)
+
+      assert_equal conversation_id, event.conversation_id
+    end
+
+    def test_attributes_constant_values_can_be_passed_as_args_and_set_on_init
+      assert_includes NewRelic::Agent::Llm::ChatCompletionMessage::ATTRIBUTES, :role
+      role = 'user'
+      event = NewRelic::Agent::Llm::ChatCompletionMessage.new(role: role)
+
+      assert_equal role, event.role
+    end
+
+    def test_args_passed_to_init_not_set_as_instance_vars_when_not_in_attributes_constant
+      event = NewRelic::Agent::Llm::ChatCompletionMessage.new(fake: 'fake')
+
+      refute_includes event.attributes, :fake
+      refute event.instance_variable_defined?(:@fake)
+    end
+
+    def test_record_creates_an_event
+      in_transaction do |txn|
+        message = NewRelic::Agent::Llm::ChatCompletionMessage.new(
+          id: 7, content: 'Red-Tailed Hawk'
+        )
+        message.sequence = 2
+        message.conversation_id = 25
+        message.request_id = '789'
+        message.response_model = 'gpt-4'
+        message.vendor = 'OpenAI'
+        message.role = 'system'
+        message.completion_id = 123
+        message.is_response = 'true'
+
+        message.record
+        _, events = NewRelic::Agent.agent.custom_event_aggregator.harvest!
+        type, attributes = events[0]
+
+        assert_equal 'LlmChatCompletionMessage', type['type']
+
+        assert_equal 7, attributes['id']
+        assert_equal 25, attributes['conversation_id']
+        assert_equal '789', attributes['request_id']
+        assert_equal txn.current_segment.guid, attributes['span_id']
+        assert_equal txn.guid, attributes['transaction_id']
+        assert_equal txn.trace_id, attributes['trace_id']
+        assert_equal 'gpt-4', attributes['response_model']
+        assert_equal 'OpenAI', attributes['vendor']
+        assert_equal 'Ruby', attributes['ingest_source']
+        assert_equal 'Red-Tailed Hawk', attributes['content']
+        assert_equal 'system', attributes['role']
+        assert_equal 2, attributes['sequence']
+        assert_equal 123, attributes['completion_id']
+        assert_equal 'true', attributes['is_response']
+      end
+    end
+  end
+end