Add support for llm. custom attributes

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

@@ -71,7 +71,8 @@ def create_chat_completion_summary(parameters)
         conversation_id: conversation_id,
         request_max_tokens: parameters[:max_tokens] || parameters['max_tokens'],
         request_model: parameters[:model] || parameters['model'],
-        temperature: parameters[:temperature] || parameters['temperature']
+        temperature: parameters[:temperature] || parameters['temperature'],
+        metadata: llm_custom_attributes
       )
     end
 
@@ -80,7 +81,8 @@ def create_embeddings_event(parameters)
         # TODO: POST-GA: Add metadata from add_custom_attributes if prefixed with 'llm.', except conversation_id
         vendor: VENDOR,
         input: parameters[:input] || parameters['input'],
-        request_model: parameters[:model] || parameters['model']
+        request_model: parameters[:model] || parameters['model'],
+        metadata: llm_custom_attributes
       )
     end
 
@@ -144,9 +146,16 @@ def update_chat_completion_messages(messages, response, summary)
         message.conversation_id = conversation_id
         message.request_id = summary.request_id
         message.response_model = response['model']
+        message.metadata = llm_custom_attributes
       end
     end
 
+    def llm_custom_attributes
+      attributes = NewRelic::Agent::Tracer.current_transaction&.attributes&.custom_attributes&.select { |k| k.to_s.match(/llm.*/) }
+
+      attributes&.transform_keys! { |key| key[4..-1] }
+    end
+
     def record_openai_metric
       NewRelic::Agent.record_metric(nr_supportability_metric, 0.0)
     end

lib/new_relic/agent/llm/llm_event.rb

@@ -9,7 +9,7 @@ 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 trace_id response_model vendor
-          ingest_source]
+          ingest_source metadata]
         # These attributes should not be passed as arguments to initialize and will be set by the agent
         AGENT_DEFINED_ATTRIBUTES = %i[span_id trace_id ingest_source]
         # Some attributes have names that can't be written as symbols used for metaprogramming.
@@ -51,9 +51,12 @@ def initialize(opts = {})
         # 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|
+          attributes_hash = attributes.each_with_object({}) do |attr, hash|
             hash[replace_attr_with_string(attr)] = instance_variable_get(:"@#{attr}")
           end
+          attributes_hash.merge!(metadata) && attributes_hash.delete(:metadata) if !metadata.nil?
+
+          attributes_hash
         end
 
         # Subclasses define an attributes method to concatenate attributes

test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb

@@ -130,10 +130,13 @@ def test_set_llm_agent_attribute_on_chat_transaction
     assert_truthy harvest_transaction_events![1][0][2][:llm]
   end
 
-  def test_conversation_id_added_to_summary_events
-    conversation_id = '12345'
+  def test_llm_custom_attributes_added_to_summary_events
     in_transaction do
-      NewRelic::Agent.add_custom_attributes({'llm.conversation_id' => conversation_id})
+      NewRelic::Agent.add_custom_attributes({
+        'llm.conversation_id' => '1993',
+        'llm.JurassicPark' => 'Steven Spielberg',
+        'trex' => 'carnivore'
+      })
       stub_post_request do
         client.chat(parameters: chat_params)
       end
@@ -142,24 +145,48 @@ def test_conversation_id_added_to_summary_events
     _, events = @aggregator.harvest!
     summary_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionSummary::EVENT_NAME }
 
-    assert_equal conversation_id, summary_event[1]['conversation_id']
+    assert_equal '1993', summary_event[1]['conversation_id']
+    assert_equal 'Steven Spielberg', summary_event[1]['JurassicPark']
+    refute summary_event[1]['trex']
   end
 
-  def test_conversation_id_added_to_message_events
-    conversation_id = '12345'
+  def test_llm_custom_attributes_added_to_embedding_events
+    in_transaction do
+      NewRelic::Agent.add_custom_attributes({
+        'llm.conversation_id' => '1997',
+        'llm.TheLostWorld' => 'Steven Spielberg',
+        'triceratops' => 'herbivore'
+      })
+      stub_post_request do
+        client.embeddings(parameters: chat_params)
+      end
+    end
+    _, events = @aggregator.harvest!
+    embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME }
+
+    assert_equal '1997', embedding_event[1]['conversation_id']
+    assert_equal 'Steven Spielberg', embedding_event[1]['TheLostWorld']
+    refute embedding_event[1]['fruit']
+  end
 
+  def test_llm_custom_attributes_added_to_message_events
     in_transaction do
-      NewRelic::Agent.add_custom_attributes({'llm.conversation_id' => conversation_id})
+      NewRelic::Agent.add_custom_attributes({
+        'llm.conversation_id' => '2001',
+        'llm.JurassicParkIII' => 'Joe Johnston',
+        'Pterosaur' => 'Can fly — scary!'
+      })
       stub_post_request do
         client.chat(parameters: chat_params)
       end
     end
-
     _, events = @aggregator.harvest!
     message_events = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME }
 
     message_events.each do |event|
-      assert_equal conversation_id, event[1]['conversation_id']
+      assert_equal '2001', event[1]['conversation_id']
+      assert_equal 'Joe Johnston', event[1]['JurassicParkIII']
+      refute event[1]['Pterosaur']
     end
   end
 

test/new_relic/agent/llm/llm_event_test.rb

@@ -46,6 +46,18 @@ def test_event_attributes_returns_a_hash_of_assigned_attributes_and_values
       assert_equal('gpt-4', result['response.model'])
     end
 
+    def test_event_attributes_adds_custom_attributes
+      event = NewRelic::Agent::Llm::LlmEvent.new(id: 123)
+      event.vendor = 'OpenAI'
+      event.response_model = 'gpt-4'
+      event.metadata = {'Marathon' => '26.2', 'Ultra Marathon' => 'Ouch'}
+      result = event.event_attributes
+
+      assert_equal('26.2', result['Marathon'])
+      assert_equal('Ouch', result['Ultra Marathon'])
+      assert_equal('OpenAI', result[:vendor])
+    end
+
     def test_record_does_not_create_an_event
       event = NewRelic::Agent::Llm::LlmEvent.new
       event.record