Create set_llm_token_count_callback API

lib/new_relic/agent.rb

@@ -109,9 +109,11 @@ class SerializationError < StandardError; end
     LLM_FEEDBACK_MESSAGE = 'LlmFeedbackMessage'
 
     attr_reader :error_group_callback
+    attr_reader :llm_token_count_callback
 
     @agent = nil
     @error_group_callback = nil
+    @llm_token_count_callback = nil
     @logger = nil
     @tracer_lock = Mutex.new
     @tracer_queue = []
@@ -436,6 +438,42 @@ def record_llm_feedback_event(trace_id:,
 
     # @!endgroup
 
+    # @!group LLM callbacks
+
+    # Set a callback proc for calculating `token_count` attributes for
+    # LlmEmbedding and LlmChatCompletionMessage events
+    #
+    # @param callback_proc [Proc] the callback proc
+    #
+    # Typically this method should be called only once to set a callback for
+    # use with all LLM token calculations. If it is called multiple times, each
+    # new callback will replace the old one.
+    #
+    # The proc will be called with a single hash as its input argument and
+    # must return an Integer representing the number of tokens used for that
+    # particular prompt, completion message, or embedding. Values less than or
+    # equal to 0 will not be attached to an event.
+    #
+    # The hash has the following keys:
+    #
+    # :model => [String] The name of the LLM model
+    # :content => [String] The message content or prompt
+    #
+    # @api public
+    #
+    def set_llm_token_count_callback(callback_proc)
+      unless callback_proc.is_a?(Proc)
+        NewRelic::Agent.logger.error("#{self}.#{__method__}: expected an argument of type Proc, " \
+                                     "got #{callback_proc.class}")
+        return
+      end
+
+      record_api_supportability_metric(:set_llm_token_count_callback)
+      @llm_token_count_callback = callback_proc
+    end
+
+    # @!endgroup
+
     # @!group Manual agent configuration and startup/shutdown
 
     # Call this to manually start the agent in situations where the agent does

lib/new_relic/supportability_helper.rb

@@ -46,6 +46,7 @@ module SupportabilityHelper
       :recording_web_transaction?,
       :require_test_helper,
       :set_error_group_callback,
+      :set_llm_token_count_callback,
       :set_segment_callback,
       :set_sql_obfuscator,
       :set_transaction_name,

test/new_relic/agent_test.rb

@@ -587,6 +587,60 @@ def test_successful_set_error_group_callback_api_invocation_produces_supportabil
       assert called
     end
 
+    def test_set_llm_token_count_callback
+      lambda = -> { 'Hello, World!' }
+      NewRelic::Agent.set_llm_token_count_callback(lambda)
+
+      assert_equal lambda,
+        NewRelic::Agent.instance_variable_get(:@llm_token_count_callback),
+        'Supplied llm token callback proc was not found to be set.'
+    end
+
+    def test_set_llm_token_count_callback_can_be_called_multiple_times
+      procs = [proc { 'one' }, proc { 'two' }, proc { 'three' }]
+      procs.each { |proc| NewRelic::Agent.set_llm_token_count_callback(proc) }
+
+      assert_equal procs.last,
+        NewRelic::Agent.instance_variable_get(:@llm_token_count_callback),
+        'Supplied llm token callback proc was not found to be set.'
+    end
+
+    def test_set_llm_token_count_callback_rejects_non_proc
+      skip_unless_minitest5_or_above
+
+      NewRelic::Agent.instance_variable_set(:@llm_token_count_callback, nil)
+
+      mock_logger = MiniTest::Mock.new
+      mock_logger.expect :error, nil, [/expected an argument of type Proc/]
+
+      NewRelic::Agent.stub(:logger, mock_logger) do
+        NewRelic::Agent.set_llm_token_count_callback([])
+      end
+
+      mock_logger.verify
+
+      assert_nil NewRelic::Agent.instance_variable_get(:@llm_token_count_callback)
+    end
+
+    def test_llm_token_count_callback_is_exposed
+      callback = 'shhhhhh'
+      NewRelic::Agent.instance_variable_set(:@llm_token_count_callback, callback)
+
+      assert_equal callback, NewRelic::Agent.llm_token_count_callback
+    ensure
+      NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback)
+    end
+
+    def test_successful_set_llm_token_count_callback_api_invocation_produces_supportability_metrics
+      called = false
+      verification_proc = proc { |name| called = true if name == :set_llm_token_count_callback }
+      NewRelic::Agent.stub :record_api_supportability_metric, verification_proc do
+        NewRelic::Agent.set_llm_token_count_callback(proc {})
+      end
+
+      assert called
+    end
+
     def test_set_user_id_attribute
       test_user = 'test_user_id'