LLM events + chat completion message

lib/new_relic/agent/llm_event.rb

@@ -3,18 +3,21 @@
 # frozen_string_literal: true
 
 require_relative 'llm_event/chat_completion'
+# require_relative 'llm_event/embedding'
+# require_relative 'llm_event/feedback'
+require_relative 'llm_event/chat_completion/message'
+# require_relative 'llm_event/chat_completion/summary'
 
 module NewRelic
   module Agent
     class LlmEvent
-
       # response_model looks like repsonse.model
       attr_accessor :id, :app_name, :request_id, :span_id, :transaction_id, :trace_id, :response_model, :vendor
       INGEST_SOURCE = 'Ruby'
 
-      def initialize(id:, app_name:, request_id:, span_id:, transaction_id:, trace_id:, response_model:, vendor:, **args)
+      def initialize(id: nil, request_id: nil, span_id: nil, transaction_id: nil, trace_id: nil, response_model: nil, vendor: nil, **args)
         @id = id
-        @app_name = app_name
+        @app_name = NewRelic::Agent.config[:app_name]
         @request_id = request_id
         @span_id = span_id
         @transaction_id = transaction_id
@@ -23,8 +26,13 @@ def initialize(id:, app_name:, request_id:, span_id:, transaction_id:, trace_id:
         @vendor = vendor
       end
 
+      def llm_event_attributes
+        {id: @id, app_name: @app_name}
+      end
+
       # Method for subclasses to override
-      def record; end
+      def record
+      end
     end
   end
-end
+end

lib/new_relic/agent/llm_event/chat_completion.rb

@@ -2,27 +2,29 @@
 # See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
 # frozen_string_literal: true
 
-
 module NewRelic
   module Agent
     class LlmEvent
-      class ChatCompletion < NewRelic::Agent::LlmEvent
-
+      class ChatCompletion < LlmEvent
         # Real metrics are written: request.max_tokens, response.number_of_messages
         attr_accessor :api_key_last_four_digits, :conversation_id, :request_max_tokens, :response_number_of_messages
 
-        def initialize(api_key_last_four_digits:, conversation_id:, request_max_tokens:, response_number_of_messages:)
+        def initialize(api_key_last_four_digits:, conversation_id:, request_max_tokens:, response_number_of_messages:, **args)
           @api_key_last_four_digits = api_key_last_four_digits
           @conversation_id = conversation_id
           @request_max_tokens = request_max_tokens
           @response_number_of_messages = response_number_of_messages
           super
         end
 
-      # Method for subclasses to override
-      def record; end
+        def chat_completion_attributes
+          {api_key_last_four_digits: @api_key_last_four_digits, conversation_id: @conversation_id}
+        end
+
+        # Method for subclasses to override
+        def record
+        end
       end
     end
   end
 end
-

lib/new_relic/agent/llm_event/chat_completion/message.rb

@@ -2,23 +2,32 @@
 # See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
 # frozen_string_literal: true
 
-class NewRelic
-  class Agent
+module NewRelic
+  module Agent
     class LlmEvent
       class ChatCompletion
-        class Message < NewRelic::Agent::LlmEvent::ChatCompletion
-
+        class Message < ChatCompletion
           attr_accessor :content, :role, :sequence, :completion_id, :is_response
+          EVENT_NAME = 'LlmChatCompletionMessage'
+
+          def initialize(content: nil, role: nil, sequence: nil, completion_id: nil, is_response: nil, **args)
+            @content = content
+            @role = role
+            @sequence = sequence
+            @completion_id = completion_id
+            @is_response = is_response
+            super
+          end
 
-          def initialize
-
+          def message_attributes
+            {content: @content, role: @role}.merge(chat_completion_attributes, llm_event_attributes)
           end
-    
+
           def record
-
+            NewRelic::Agent.record_custom_event(EVENT_NAME, message_attributes)
           end
         end
       end
     end
   end
-end
+end

test/new_relic/agent/llm_event_test.rb

@@ -3,17 +3,35 @@
 # frozen_string_literal: true
 
 require_relative '../../test_helper'
-
+require_relative '../../agent_helper'
 
 module NewRelic::Agent
   class LlmEventTest < Minitest::Test
     # def test_attributes
     #   NewRelic::Agent::LlmEvent.new
-
-    # end
+    def setup
+      events = NewRelic::Agent.instance.events
+      @aggregator = NewRelic::Agent::CustomEventAggregator.new(events)
+    end
 
     def test_attributes_chat
-      NewRelic::Agent::LlmEvent::ChatCompletion.new
+      NewRelic::Agent::LlmEvent::ChatCompletion::Message.new(id: 123)
+    end
+
+    def test_attribute_merge
+      message = NewRelic::Agent::LlmEvent::ChatCompletion::Message.new(content: 'hi', role: 'speaker', api_key_last_four_digits: 'sk-0', conversation_id: 123, id: 345, app_name: NewRelic::Agent.config[:app_name])
+      message.record
+      binding.irb
+      _, events = @aggregator.harvest!
+      # NewRelic::Agent.agent.send(:harvest_and_send_custom_event_data)
+      # returned = first_call_for('custom_event_data').events
+      # events.first[0].delete('priority')
+      # event = events.first
+
+      expected_event = [{'type' => 'DummyType', 'timestamp' => 'bhjbjh'},
+        {'foo' => 'bar', 'baz' => 'qux'}]
+
+      assert_equal(expected_event, events)
     end
   end
-end
+end