From b5724433e1a847220b7a5e0b1a740df093219650 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Tue, 5 Mar 2024 14:37:28 -0800 Subject: [PATCH 01/13] Add token_count attribute This attribute is available on ChatCompletionMessage and Embedding --- lib/new_relic/agent/llm/chat_completion_message.rb | 3 ++- lib/new_relic/agent/llm/embedding.rb | 2 +- test/new_relic/agent/llm/chat_completion_message_test.rb | 2 ++ test/new_relic/agent/llm/embedding_test.rb | 2 ++ 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/lib/new_relic/agent/llm/chat_completion_message.rb b/lib/new_relic/agent/llm/chat_completion_message.rb index eca258e1ad..33d6ca9c47 100644 --- a/lib/new_relic/agent/llm/chat_completion_message.rb +++ b/lib/new_relic/agent/llm/chat_completion_message.rb @@ -8,7 +8,8 @@ module Llm class ChatCompletionMessage < LlmEvent include ChatCompletion - ATTRIBUTES = %i[content role sequence completion_id is_response] + ATTRIBUTES = %i[content role sequence completion_id token_count + is_response] EVENT_NAME = 'LlmChatCompletionMessage' attr_accessor(*ATTRIBUTES) diff --git a/lib/new_relic/agent/llm/embedding.rb b/lib/new_relic/agent/llm/embedding.rb index b458d30a19..8e23d6b15f 100644 --- a/lib/new_relic/agent/llm/embedding.rb +++ b/lib/new_relic/agent/llm/embedding.rb @@ -8,7 +8,7 @@ module Llm class Embedding < LlmEvent include ResponseHeaders - ATTRIBUTES = %i[input request_model duration error] + ATTRIBUTES = %i[input request_model token_count duration error] ATTRIBUTE_NAME_EXCEPTIONS = { request_model: 'request.model' } diff --git a/test/new_relic/agent/llm/chat_completion_message_test.rb b/test/new_relic/agent/llm/chat_completion_message_test.rb index 9a23530363..8087e4c55c 100644 --- a/test/new_relic/agent/llm/chat_completion_message_test.rb +++ b/test/new_relic/agent/llm/chat_completion_message_test.rb @@ -67,6 +67,7 @@ def test_record_creates_an_event message.role = 'system' message.completion_id = 123 message.is_response = 'true' + message.token_count = 10 message.record _, events = NewRelic::Agent.agent.custom_event_aggregator.harvest! @@ -87,6 +88,7 @@ def test_record_creates_an_event assert_equal 2, attributes['sequence'] assert_equal 123, attributes['completion_id'] assert_equal 'true', attributes['is_response'] + assert_equal 10, attributes['token_count'] end end end diff --git a/test/new_relic/agent/llm/embedding_test.rb b/test/new_relic/agent/llm/embedding_test.rb index cf809fe208..8a9a61842a 100644 --- a/test/new_relic/agent/llm/embedding_test.rb +++ b/test/new_relic/agent/llm/embedding_test.rb @@ -53,6 +53,7 @@ def test_record_creates_an_event embedding.vendor = 'OpenAI' embedding.duration = '500' embedding.error = 'true' + embedding.token_count = 10 embedding.llm_version = '2022-01-01' embedding.rate_limit_requests = '100' embedding.rate_limit_tokens = '101' @@ -79,6 +80,7 @@ def test_record_creates_an_event assert_equal 'Ruby', attributes['ingest_source'] assert_equal '500', attributes['duration'] assert_equal 'true', attributes['error'] + assert_equal 10, attributes['token_count'] assert_equal '2022-01-01', attributes['response.headers.llm_version'] assert_equal '100', attributes['response.headers.ratelimitLimitRequests'] assert_equal '101', attributes['response.headers.ratelimitLimitTokens'] From 99f874b910ecc4f665525e8b187c2eb3935a0d9f Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Tue, 5 Mar 2024 15:38:49 -0800 Subject: [PATCH 02/13] WIP token count assignment --- .../ruby_openai/instrumentation.rb | 25 +++++++++++++++-- .../suites/ruby_openai/openai_helpers.rb | 9 +++++- .../ruby_openai_instrumentation_test.rb | 28 ++++++++++++++++++- 3 files changed, 58 insertions(+), 4 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index 678f0a24c0..1b960e1230 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -93,6 +93,7 @@ def add_chat_completion_response_params(parameters, response, event) def add_embeddings_response_params(response, event) event.response_model = response['model'] + event.token_count = response.dig('usage', 'prompt_tokens') || NewRelic::Agent.llm_token_count_callback.call({model: event.response_model, content: event.input}) end # The customer must call add_custom_attributes with llm.conversation_id @@ -110,8 +111,7 @@ def create_chat_completion_messages(parameters, summary_id) role: message[:role] || message['role'], sequence: index, completion_id: summary_id, - vendor: VENDOR, - is_response: true + vendor: VENDOR ) end end @@ -139,9 +139,22 @@ 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.token_count = calculate_message_token_count(message, response) end end + def calculate_message_token_count(message, response) + token_count = if message.is_response + response.dig('usage', 'completion_tokens') + else + response.dig('usage', 'prompt_tokens') + end + + return NewRelic::Agent.llm_token_count_callback.call({model: response['model'], content: message.content}) if token_count.nil? + + token_count + end + def record_openai_metric NewRelic::Agent.record_metric(nr_supportability_metric, 0.0) end @@ -154,6 +167,14 @@ def nr_supportability_metric @nr_supportability_metric ||= "Supportability/Ruby/ML/OpenAI/#{::OpenAI::VERSION}" end + def llm_token_count_callback + NewRelic::Agent.llm_token_count_callback + end + + def build_llm_token_count_callback_hash(model, content) + { model: model, content: content } + end + def finish(segment, event) segment&.finish diff --git a/test/multiverse/suites/ruby_openai/openai_helpers.rb b/test/multiverse/suites/ruby_openai/openai_helpers.rb index 8bd0e3643b..c25fd22c50 100644 --- a/test/multiverse/suites/ruby_openai/openai_helpers.rb +++ b/test/multiverse/suites/ruby_openai/openai_helpers.rb @@ -124,6 +124,13 @@ def faraday_connection.post(*args); ChatResponse.new; end faraday_connection end + def embedding_faraday_connection + faraday_connection = Faraday.new + def faraday_connection.post(*args); EmbeddingsResponse.new; end + + faraday_connection + end + def error_faraday_connection faraday_connection = Faraday.new def faraday_connection.post(*args); raise 'deception'; end @@ -201,7 +208,7 @@ def stub_embeddings_post_request(&blk) yield end else - connection_client.stub(:conn, faraday_connection) do + connection_client.stub(:conn, embedding_faraday_connection) do yield end end diff --git a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb index 0dff89aaf9..601afcc49d 100644 --- a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb +++ b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb @@ -225,7 +225,7 @@ def test_embedding_event_sets_error_true_if_raised end def test_set_llm_agent_attribute_on_embedding_transaction - in_transaction do |txn| + in_transaction do stub_embeddings_post_request do client.embeddings(parameters: embeddings_params) end @@ -233,4 +233,30 @@ def test_set_llm_agent_attribute_on_embedding_transaction assert_truthy harvest_transaction_events![1][0][2][:llm] end + + def test_token_count_recorded_from_usage_object_when_present_on_embeddings + in_transaction do + stub_embeddings_post_request do + client.embeddings(parameters: embeddings_params) + end + end + + _, events = @aggregator.harvest! + embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } + + assert_equal EmbeddingsResponse.new.body['usage']['prompt_tokens'], embedding_event[1]['token_count'] + end + + def test_token_count_recorded_from_callback_when_usage_is_missing_on_embeddings + # how to stub the missing usage?? + end + + def test_token_count_recorded_when_message_not_response_and_usage_present + end + + def test_token_count_recorded_when_message_is_response_and_usage_present + end + + def test_token_count_recorded_from_callback_when_token_count_nil + end end From 138d36087dfad551dcc8cad195bcd1f46241f1db Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 6 Mar 2024 14:49:57 -0800 Subject: [PATCH 03/13] WIP more token count assignment updates --- .../ruby_openai/instrumentation.rb | 32 +++++++--- .../ruby_openai_instrumentation_test.rb | 63 +++++++++++++++++-- 2 files changed, 83 insertions(+), 12 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index 1b960e1230..648dd9dcbd 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -54,7 +54,7 @@ def chat_completions_instrumentation(parameters) # TODO: Remove !response.include?('error') when we drop support for versions below 4.0.0 if response && !response.include?('error') add_chat_completion_response_params(parameters, response, event) - messages = update_chat_completion_messages(messages, response, event) + messages = update_chat_completion_messages(messages, response, event, parameters) end response @@ -93,7 +93,7 @@ def add_chat_completion_response_params(parameters, response, event) def add_embeddings_response_params(response, event) event.response_model = response['model'] - event.token_count = response.dig('usage', 'prompt_tokens') || NewRelic::Agent.llm_token_count_callback.call({model: event.response_model, content: event.input}) + event.token_count = response.dig('usage', 'prompt_tokens') || NewRelic::Agent.llm_token_count_callback&.call({model: event.request_model, content: event.input}) end # The customer must call add_custom_attributes with llm.conversation_id @@ -129,7 +129,7 @@ def create_chat_completion_response_messages(response, sequence_origin, summary_ end end - def update_chat_completion_messages(messages, response, summary) + def update_chat_completion_messages(messages, response, summary, parameters) messages += create_chat_completion_response_messages(response, messages.size, summary.id) response_id = response['id'] || NewRelic::Agent::GuidGenerator.generate_guid @@ -139,20 +139,38 @@ 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.token_count = calculate_message_token_count(message, response) + message.token_count = calculate_message_token_count(message, response, parameters) end end - def calculate_message_token_count(message, response) + def calculate_message_token_count(message, response, parameters) + # message is response + # more than one message in response + # use the callback + # one message in response + # use the usage object + # message is request + # more than one message in request + # use the callback + # one message in request + # use the usage object + + request_message_length = (parameters['messages']|| parameters[:messages]).length + + response_message_length = response['choices'].length + binding.irb + + return NewRelic::Agent.llm_token_count_callback&.call({ model: response['model'], content: message.content }) unless message.is_response && (request_message_length > 1) + token_count = if message.is_response response.dig('usage', 'completion_tokens') else response.dig('usage', 'prompt_tokens') end - return NewRelic::Agent.llm_token_count_callback.call({model: response['model'], content: message.content}) if token_count.nil? + return NewRelic::Agent.llm_token_count_callback&.call({model: response['model'], content: message.content}) if token_count.nil? - token_count + token_count.to_i end def record_openai_metric diff --git a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb index 601afcc49d..4e375ed7f5 100644 --- a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb +++ b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb @@ -6,6 +6,9 @@ class RubyOpenAIInstrumentationTest < Minitest::Test include OpenAIHelpers + # some of the private methods are too difficult to stub + # we can test them directly by including the module + include NewRelic::Agent::Instrumentation::OpenAI def setup @aggregator = NewRelic::Agent.agent.custom_event_aggregator @@ -247,16 +250,66 @@ def test_token_count_recorded_from_usage_object_when_present_on_embeddings assert_equal EmbeddingsResponse.new.body['usage']['prompt_tokens'], embedding_event[1]['token_count'] end - def test_token_count_recorded_from_callback_when_usage_is_missing_on_embeddings - # how to stub the missing usage?? + def test_token_count_nil_when_usage_is_missing_on_embeddings_and_no_callback_defined + mock_response = {'model' => 'gpt-2001'} + mock_event = NewRelic::Agent::Llm::Embedding.new(request_model: 'gpt-2004', input: 'what does my dog want?') + add_embeddings_response_params(mock_response, mock_event) + + assert_nil mock_event.token_count + end + + def test_token_count_assigned_by_callback_when_usage_is_missing_and_callback_defined + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 7734 }) + + mock_response = {'model' => 'gpt-2001'} + mock_event = NewRelic::Agent::Llm::Embedding.new(request_model: 'gpt-2004', input: 'what does my dog want?') + add_embeddings_response_params(mock_response, mock_event) + + assert_equal 7734, mock_event.token_count + + NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) + end + + def test_token_count_when_message_not_response_and_usage_present_and_only_one_request_message + message = NewRelic::Agent::Llm::ChatCompletionMessage.new(content: 'pineapple strawberry') + response = {'usage' => {'prompt_tokens' => 123456, 'completion_tokens' => 654321}, 'model' => 'gpt-2001'} + parameters = {'messages' => ['one']} + + result = calculate_message_token_count(message, response, parameters) + + assert_equal 123456, result + end + + def test_token_count_when_message_not_response_and_usage_present_and_only_one_request_message_and_messages_params_symbol + message = NewRelic::Agent::Llm::ChatCompletionMessage.new(content: 'pineapple strawberry') + response = {'usage' => {'prompt_tokens' => 123456, 'completion_tokens' => 654321}, 'model' => 'gpt-2001'} + parameters = {:messages => ['one']} + + calculate_message_token_count(message, response, parameters) + + assert_equal 123456, message.token_count end - def test_token_count_recorded_when_message_not_response_and_usage_present + def test_token_count_when_message_not_response_and_usage_present_and_multiple_request_messages_but_no_callback + in_transaction do + stub_post_request do + result = client.chat(parameters: chat_params) + end + end + + _, events = @aggregator.harvest! + chat_completion_messages = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + + not_response_messages = chat_completion_messages.find { |event| !event[1].key?('is_response') } + + not_response_messages.each do |msg| + assert_nil msg.token_count + end end - def test_token_count_recorded_when_message_is_response_and_usage_present + def test_token_count_when_message_is_response_and_usage_present end - def test_token_count_recorded_from_callback_when_token_count_nil + def test_token_count_from_callback_when_token_count_nil end end From 46b5f2de187247ee705e838d141cec7687603597 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 6 Mar 2024 14:56:03 -0800 Subject: [PATCH 04/13] More cleanup --- .../agent/instrumentation/ruby_openai/instrumentation.rb | 1 - test/new_relic/agent/llm/chat_completion_message_test.rb | 2 -- 2 files changed, 3 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index 0d1ce565c2..d6c411215c 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -148,7 +148,6 @@ def calculate_message_token_count(message, response, parameters) request_message_length = (parameters['messages']|| parameters[:messages]).length response_message_length = response['choices'].length - binding.irb return NewRelic::Agent.llm_token_count_callback&.call({ model: response['model'], content: message.content }) unless message.is_response && (request_message_length > 1) diff --git a/test/new_relic/agent/llm/chat_completion_message_test.rb b/test/new_relic/agent/llm/chat_completion_message_test.rb index 443a8d82f8..85eb809982 100644 --- a/test/new_relic/agent/llm/chat_completion_message_test.rb +++ b/test/new_relic/agent/llm/chat_completion_message_test.rb @@ -54,7 +54,6 @@ def test_record_creates_an_event message.response_model = 'gpt-4' message.vendor = 'OpenAI' message.role = 'system' - message.completion_id = 123 message.is_response = 'true' message.token_count = 10 @@ -74,7 +73,6 @@ def test_record_creates_an_event 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'] assert_equal 10, attributes['token_count'] end From 45513189cebd9d4818f431572462efb2da53847f Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 6 Mar 2024 15:03:47 -0800 Subject: [PATCH 05/13] Rubocop --- .../ruby_openai/instrumentation.rb | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index d6c411215c..a8a69a53af 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -135,21 +135,21 @@ def update_chat_completion_messages(messages, response, summary, parameters) def calculate_message_token_count(message, response, parameters) # message is response - # more than one message in response - # use the callback - # one message in response - # use the usage object + # more than one message in response + # use the callback + # one message in response + # use the usage object # message is request - # more than one message in request - # use the callback - # one message in request - # use the usage object + # more than one message in request + # use the callback + # one message in request + # use the usage object - request_message_length = (parameters['messages']|| parameters[:messages]).length + request_message_length = (parameters['messages'] || parameters[:messages]).length - response_message_length = response['choices'].length + response_message_length = response['choices'].length - return NewRelic::Agent.llm_token_count_callback&.call({ model: response['model'], content: message.content }) unless message.is_response && (request_message_length > 1) + return NewRelic::Agent.llm_token_count_callback&.call({model: response['model'], content: message.content}) unless message.is_response && (request_message_length > 1) token_count = if message.is_response response.dig('usage', 'completion_tokens') @@ -185,7 +185,7 @@ def llm_token_count_callback end def build_llm_token_count_callback_hash(model, content) - { model: model, content: content } + {model: model, content: content} end def finish(segment, event) From 8df6886eed18d119c517f7169694885a2475d6d0 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Fri, 8 Mar 2024 11:52:09 -0800 Subject: [PATCH 06/13] WIP use just the callback, no request/response --- .../ruby_openai/instrumentation.rb | 48 +++-------- .../ruby_openai_instrumentation_test.rb | 84 ++++++++++--------- 2 files changed, 56 insertions(+), 76 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index a8a69a53af..c5e922bd6e 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -54,7 +54,7 @@ def chat_completions_instrumentation(parameters) # TODO: Remove !response.include?('error') when we drop support for versions below 4.0.0 if response && !response.include?('error') add_chat_completion_response_params(parameters, response, event) - messages = update_chat_completion_messages(messages, response, event, parameters) + messages = update_chat_completion_messages(messages, response, event) end response @@ -92,7 +92,7 @@ def add_chat_completion_response_params(parameters, response, event) def add_embeddings_response_params(response, event) event.response_model = response['model'] - event.token_count = response.dig('usage', 'prompt_tokens') || NewRelic::Agent.llm_token_count_callback&.call({model: event.request_model, content: event.input}) + event.token_count = calculate_token_count(event.request_model, event.input) end def create_chat_completion_messages(parameters, summary_id) @@ -120,7 +120,7 @@ def create_chat_completion_response_messages(response, sequence_origin, summary_ end end - def update_chat_completion_messages(messages, response, summary, parameters) + 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 @@ -129,37 +129,19 @@ def update_chat_completion_messages(messages, response, summary, parameters) message.request_id = summary.request_id message.response_model = response['model'] message.metadata = llm_custom_attributes - message.token_count = calculate_message_token_count(message, response, parameters) - end - end - - def calculate_message_token_count(message, response, parameters) - # message is response - # more than one message in response - # use the callback - # one message in response - # use the usage object - # message is request - # more than one message in request - # use the callback - # one message in request - # use the usage object - - request_message_length = (parameters['messages'] || parameters[:messages]).length - - response_message_length = response['choices'].length - return NewRelic::Agent.llm_token_count_callback&.call({model: response['model'], content: message.content}) unless message.is_response && (request_message_length > 1) + model = message.is_response ? message.response_model : summary.request_model - token_count = if message.is_response - response.dig('usage', 'completion_tokens') - else - response.dig('usage', 'prompt_tokens') + message.token_count = calculate_token_count(model, message.content) end + end - return NewRelic::Agent.llm_token_count_callback&.call({model: response['model'], content: message.content}) if token_count.nil? + def calculate_token_count(model, content) + # return unless NewRelic::Agent.config['ai_monitoring.record_content.enabled'] + return unless NewRelic::Agent.llm_token_count_callback - token_count.to_i + count = NewRelic::Agent.llm_token_count_callback.call({model: model, content: content}) + return count unless count.is_a?(Integer) && count <= 0 end def llm_custom_attributes @@ -180,14 +162,6 @@ def nr_supportability_metric @nr_supportability_metric ||= "Supportability/Ruby/ML/OpenAI/#{::OpenAI::VERSION}" end - def llm_token_count_callback - NewRelic::Agent.llm_token_count_callback - end - - def build_llm_token_count_callback_hash(model, content) - {model: model, content: content} - end - def finish(segment, event) segment&.finish diff --git a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb index 09f768b379..c6b3ab226f 100644 --- a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb +++ b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb @@ -246,7 +246,9 @@ def test_set_llm_agent_attribute_on_embedding_transaction assert_truthy harvest_transaction_events![1][0][2][:llm] end - def test_token_count_recorded_from_usage_object_when_present_on_embeddings + def test_embeddings_token_count_assigned_by_callback_if_present + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 7734 }) + in_transaction do stub_embeddings_post_request do client.embeddings(parameters: embeddings_params) @@ -256,69 +258,73 @@ def test_token_count_recorded_from_usage_object_when_present_on_embeddings _, events = @aggregator.harvest! embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } - assert_equal EmbeddingsResponse.new.body['usage']['prompt_tokens'], embedding_event[1]['token_count'] - end + assert_equal 7734, embedding_event[1]['token_count'] - def test_token_count_nil_when_usage_is_missing_on_embeddings_and_no_callback_defined - mock_response = {'model' => 'gpt-2001'} - mock_event = NewRelic::Agent::Llm::Embedding.new(request_model: 'gpt-2004', input: 'what does my dog want?') - add_embeddings_response_params(mock_response, mock_event) - - assert_nil mock_event.token_count + NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end - def test_token_count_assigned_by_callback_when_usage_is_missing_and_callback_defined - NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 7734 }) + def test_embeddings_token_count_attribute_absent_if_callback_returns_nil + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| nil }) - mock_response = {'model' => 'gpt-2001'} - mock_event = NewRelic::Agent::Llm::Embedding.new(request_model: 'gpt-2004', input: 'what does my dog want?') - add_embeddings_response_params(mock_response, mock_event) + in_transaction do + stub_embeddings_post_request do + client.embeddings(parameters: embeddings_params) + end + end - assert_equal 7734, mock_event.token_count + _, events = @aggregator.harvest! + embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } + + refute embedding_event[1].key?('token_count') NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end - def test_token_count_when_message_not_response_and_usage_present_and_only_one_request_message - message = NewRelic::Agent::Llm::ChatCompletionMessage.new(content: 'pineapple strawberry') - response = {'usage' => {'prompt_tokens' => 123456, 'completion_tokens' => 654321}, 'model' => 'gpt-2001'} - parameters = {'messages' => ['one']} - - result = calculate_message_token_count(message, response, parameters) + def test_embeddings_token_count_attribute_absent_if_callback_returns_zero + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 0 }) - assert_equal 123456, result - end + in_transaction do + stub_embeddings_post_request do + client.embeddings(parameters: embeddings_params) + end + end - def test_token_count_when_message_not_response_and_usage_present_and_only_one_request_message_and_messages_params_symbol - message = NewRelic::Agent::Llm::ChatCompletionMessage.new(content: 'pineapple strawberry') - response = {'usage' => {'prompt_tokens' => 123456, 'completion_tokens' => 654321}, 'model' => 'gpt-2001'} - parameters = {:messages => ['one']} + _, events = @aggregator.harvest! + embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } - calculate_message_token_count(message, response, parameters) + refute embedding_event[1].key?('token_count') - assert_equal 123456, message.token_count + NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end - def test_token_count_when_message_not_response_and_usage_present_and_multiple_request_messages_but_no_callback + def test_embeddings_token_count_attribute_absent_if_no_callback_available + assert_nil NewRelic::Agent.llm_token_count_callback + in_transaction do - stub_post_request do - result = client.chat(parameters: chat_params) + stub_embeddings_post_request do + client.embeddings(parameters: embeddings_params) end end _, events = @aggregator.harvest! - chat_completion_messages = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } - not_response_messages = chat_completion_messages.find { |event| !event[1].key?('is_response') } + refute embedding_event[1].key?('token_count') + end + + def test_chat_completion_message_token_count_assigned_by_callback_if_present - not_response_messages.each do |msg| - assert_nil msg.token_count - end end - def test_token_count_when_message_is_response_and_usage_present + def test_chat_completion_message_token_count_attribute_absent_if_callback_returns_nil + end - def test_token_count_from_callback_when_token_count_nil + def test_chat_completion_message_token_count_attribute_absent_if_callback_returns_zero + + end + + def test_chat_completion_message_token_count_attribute_absent_if_no_callback_available + end end From 416a20619668afaeb5b5957f199ec99782789d74 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Fri, 8 Mar 2024 16:35:18 -0800 Subject: [PATCH 07/13] fix tests --- .../ruby_openai/instrumentation.rb | 2 +- .../ruby_openai_instrumentation_test.rb | 65 ++++++++++++++++--- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index c5e922bd6e..4f3b8e9049 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -141,7 +141,7 @@ def calculate_token_count(model, content) return unless NewRelic::Agent.llm_token_count_callback count = NewRelic::Agent.llm_token_count_callback.call({model: model, content: content}) - return count unless count.is_a?(Integer) && count <= 0 + return count if count.is_a?(Integer) && count > 0 end def llm_custom_attributes diff --git a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb index c6b3ab226f..27c7d6db8d 100644 --- a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb +++ b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb @@ -6,12 +6,9 @@ class RubyOpenAIInstrumentationTest < Minitest::Test include OpenAIHelpers - # some of the private methods are too difficult to stub - # we can test them directly by including the module - include NewRelic::Agent::Instrumentation::OpenAI - def setup @aggregator = NewRelic::Agent.agent.custom_event_aggregator + NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) if NewRelic::Agent.instance_variable_defined?(:@llm_token_count_callback) end def teardown @@ -259,8 +256,6 @@ def test_embeddings_token_count_assigned_by_callback_if_present embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } assert_equal 7734, embedding_event[1]['token_count'] - - NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end def test_embeddings_token_count_attribute_absent_if_callback_returns_nil @@ -276,8 +271,6 @@ def test_embeddings_token_count_attribute_absent_if_callback_returns_nil embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } refute embedding_event[1].key?('token_count') - - NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end def test_embeddings_token_count_attribute_absent_if_callback_returns_zero @@ -293,8 +286,6 @@ def test_embeddings_token_count_attribute_absent_if_callback_returns_zero embedding_event = events.find { |event| event[0]['type'] == NewRelic::Agent::Llm::Embedding::EVENT_NAME } refute embedding_event[1].key?('token_count') - - NewRelic::Agent.remove_instance_variable(:@llm_token_count_callback) end def test_embeddings_token_count_attribute_absent_if_no_callback_available @@ -313,18 +304,72 @@ def test_embeddings_token_count_attribute_absent_if_no_callback_available end def test_chat_completion_message_token_count_assigned_by_callback_if_present + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 7734 }) + + in_transaction do + stub_post_request do + client.chat(parameters: chat_params) + end + end + _, events = @aggregator.harvest! + messages = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + + messages.each do |message| + assert_equal 7734, message[1]['token_count'] + end end def test_chat_completion_message_token_count_attribute_absent_if_callback_returns_nil + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| nil }) + + in_transaction do + stub_post_request do + client.chat(parameters: chat_params) + end + end + + _, events = @aggregator.harvest! + messages = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + messages.each do |message| + refute message[1].key?('token_count') + end end def test_chat_completion_message_token_count_attribute_absent_if_callback_returns_zero + NewRelic::Agent.set_llm_token_count_callback(proc { |hash| 0 }) + in_transaction do + stub_post_request do + client.chat(parameters: chat_params) + end + end + + _, events = @aggregator.harvest! + messages = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + + messages.each do |message| + refute message[1].key?('token_count') + end end def test_chat_completion_message_token_count_attribute_absent_if_no_callback_available + assert_nil NewRelic::Agent.llm_token_count_callback + in_transaction do + stub_post_request do + client.chat(parameters: chat_params) + end + end + + _, events = @aggregator.harvest! + messages = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } + + refute_empty messages + + messages.each do |message| + refute message[1].key?('token_count') + end end end From 82e96a0da1e0c8752d6fa10e5cc19a1e83e1dcef Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Fri, 8 Mar 2024 16:56:24 -0800 Subject: [PATCH 08/13] Add JSON string to EmbeddingsResponse body --- .../suites/ruby_openai/openai_helpers.rb | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/test/multiverse/suites/ruby_openai/openai_helpers.rb b/test/multiverse/suites/ruby_openai/openai_helpers.rb index c25fd22c50..200bb7ac0c 100644 --- a/test/multiverse/suites/ruby_openai/openai_helpers.rb +++ b/test/multiverse/suites/ruby_openai/openai_helpers.rb @@ -33,14 +33,19 @@ def error_response(return_value: false) class EmbeddingsResponse def body(return_value: false) - {'object' => 'list', - 'data' => [{ - 'object' => 'embedding', - 'index' => 0, - 'embedding' => [0.002297497, 1, -0.016932933, 0.018126108, -0.014432343, -0.0030051514] # A real embeddings response includes dozens more vector points. - }], - 'model' => 'text-embedding-ada-002', - 'usage' => {'prompt_tokens' => 8, 'total_tokens' => 8}} + if Gem::Version.new(::OpenAI::VERSION) >= Gem::Version.new('6.0.0') || return_value + {'object' => 'list', + 'data' => [{ + 'object' => 'embedding', + 'index' => 0, + 'embedding' => [0.002297497, 1, -0.016932933, 0.018126108, -0.014432343, -0.0030051514] # A real embeddings response includes dozens more vector points. + }], + 'model' => 'text-embedding-ada-002', + 'usage' => {'prompt_tokens' => 8, 'total_tokens' => 8} + } + else + "{\"object\":\"list\",\"data\":[{\"object\":\"embedding\",\"index\":0,\"embedding\":[0.002297497,1,-0.016932933,0.018126108,-0.014432343,-0.0030051514]}],\"model\":\"text-embedding-ada-002\",\"usage\":{\"prompt_tokens\":8,\"total_tokens\":8}}" + end end end From c830e5cb2147e33b80f5571adbdbe93a98f590f3 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Fri, 8 Mar 2024 17:06:19 -0800 Subject: [PATCH 09/13] Rubocop --- test/multiverse/suites/ruby_openai/openai_helpers.rb | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/test/multiverse/suites/ruby_openai/openai_helpers.rb b/test/multiverse/suites/ruby_openai/openai_helpers.rb index 200bb7ac0c..fdcf188c43 100644 --- a/test/multiverse/suites/ruby_openai/openai_helpers.rb +++ b/test/multiverse/suites/ruby_openai/openai_helpers.rb @@ -41,10 +41,9 @@ def body(return_value: false) 'embedding' => [0.002297497, 1, -0.016932933, 0.018126108, -0.014432343, -0.0030051514] # A real embeddings response includes dozens more vector points. }], 'model' => 'text-embedding-ada-002', - 'usage' => {'prompt_tokens' => 8, 'total_tokens' => 8} - } + 'usage' => {'prompt_tokens' => 8, 'total_tokens' => 8}} else - "{\"object\":\"list\",\"data\":[{\"object\":\"embedding\",\"index\":0,\"embedding\":[0.002297497,1,-0.016932933,0.018126108,-0.014432343,-0.0030051514]}],\"model\":\"text-embedding-ada-002\",\"usage\":{\"prompt_tokens\":8,\"total_tokens\":8}}" + '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.002297497,1,-0.016932933,0.018126108,-0.014432343,-0.0030051514]}],"model":"text-embedding-ada-002","usage":{"prompt_tokens":8,"total_tokens":8}}' end end end From 131abef2070542f9eedcc24a806a2d4549fd59a9 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Mon, 11 Mar 2024 09:36:15 -0700 Subject: [PATCH 10/13] Remove unnecessary assertion --- .../suites/ruby_openai/ruby_openai_instrumentation_test.rb | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb index 7a4e446848..305d522c84 100644 --- a/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb +++ b/test/multiverse/suites/ruby_openai/ruby_openai_instrumentation_test.rb @@ -366,8 +366,6 @@ def test_chat_completion_message_token_count_attribute_absent_if_no_callback_ava _, events = @aggregator.harvest! messages = events.filter { |event| event[0]['type'] == NewRelic::Agent::Llm::ChatCompletionMessage::EVENT_NAME } - refute_empty messages - messages.each do |message| refute message[1].key?('token_count') end From 3ce45f9882a99ee600561c855285a29a63b07c67 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Mon, 11 Mar 2024 10:12:58 -0700 Subject: [PATCH 11/13] Remove extraneous return --- .../agent/instrumentation/ruby_openai/instrumentation.rb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index dfdad900ba..39565f9daf 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -147,7 +147,7 @@ def calculate_token_count(model, content) return unless NewRelic::Agent.llm_token_count_callback count = NewRelic::Agent.llm_token_count_callback.call({model: model, content: content}) - return count if count.is_a?(Integer) && count > 0 + count if count.is_a?(Integer) && count > 0 end def record_content_enabled? From e2781e373127438c7187ee9442499f2b22542d7b Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Tue, 12 Mar 2024 21:25:37 -0700 Subject: [PATCH 12/13] Remove unncessary content config check --- .../agent/instrumentation/ruby_openai/instrumentation.rb | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb index 39565f9daf..e1151e2041 100644 --- a/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb +++ b/lib/new_relic/agent/instrumentation/ruby_openai/instrumentation.rb @@ -143,7 +143,6 @@ def update_chat_completion_messages(messages, response, summary) end def calculate_token_count(model, content) - return unless record_content_enabled? return unless NewRelic::Agent.llm_token_count_callback count = NewRelic::Agent.llm_token_count_callback.call({model: model, content: content}) From cceac8e81066d39f96b26ee3d184342f5a0923e1 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 13 Mar 2024 15:51:24 -0700 Subject: [PATCH 13/13] Revert completion_id changes --- lib/new_relic/agent/llm/chat_completion_message.rb | 3 ++- test/new_relic/agent/llm/chat_completion_message_test.rb | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/lib/new_relic/agent/llm/chat_completion_message.rb b/lib/new_relic/agent/llm/chat_completion_message.rb index c678be9565..a60e654e55 100644 --- a/lib/new_relic/agent/llm/chat_completion_message.rb +++ b/lib/new_relic/agent/llm/chat_completion_message.rb @@ -6,7 +6,8 @@ module NewRelic module Agent module Llm class ChatCompletionMessage < LlmEvent - ATTRIBUTES = %i[content role sequence token_count is_response] + ATTRIBUTES = %i[content role sequence completion_id token_count + is_response] EVENT_NAME = 'LlmChatCompletionMessage' attr_accessor(*ATTRIBUTES) diff --git a/test/new_relic/agent/llm/chat_completion_message_test.rb b/test/new_relic/agent/llm/chat_completion_message_test.rb index 85eb809982..443a8d82f8 100644 --- a/test/new_relic/agent/llm/chat_completion_message_test.rb +++ b/test/new_relic/agent/llm/chat_completion_message_test.rb @@ -54,6 +54,7 @@ def test_record_creates_an_event message.response_model = 'gpt-4' message.vendor = 'OpenAI' message.role = 'system' + message.completion_id = 123 message.is_response = 'true' message.token_count = 10 @@ -73,6 +74,7 @@ def test_record_creates_an_event 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'] assert_equal 10, attributes['token_count'] end