diff --git a/newrelic/agent.py b/newrelic/agent.py index 2c7f0fb858..b433f8e317 100644 --- a/newrelic/agent.py +++ b/newrelic/agent.py @@ -153,6 +153,7 @@ def __asgi_application(*args, **kwargs): from newrelic.api.message_transaction import ( wrap_message_transaction as __wrap_message_transaction, ) +from newrelic.api.ml_model import get_ai_message_ids as __get_ai_message_ids from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper from newrelic.api.profile_trace import profile_trace as __profile_trace @@ -340,3 +341,4 @@ def __asgi_application(*args, **kwargs): insert_html_snippet = __wrap_api_call(__insert_html_snippet, "insert_html_snippet") verify_body_exists = __wrap_api_call(__verify_body_exists, "verify_body_exists") wrap_mlmodel = __wrap_api_call(__wrap_mlmodel, "wrap_mlmodel") +get_ai_message_ids = __wrap_api_call(__get_ai_message_ids, "get_ai_message_ids") diff --git a/newrelic/api/ml_model.py b/newrelic/api/ml_model.py index edbcaf3406..91f3656849 100644 --- a/newrelic/api/ml_model.py +++ b/newrelic/api/ml_model.py @@ -13,7 +13,9 @@ # limitations under the License. import sys +import warnings +from newrelic.api.transaction import current_transaction from newrelic.common.object_names import callable_name from newrelic.hooks.mlmodel_sklearn import _nr_instrument_model @@ -33,3 +35,20 @@ def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names model._nr_wrapped_label_names = label_names if metadata: model._nr_wrapped_metadata = metadata + + +def get_ai_message_ids(response_id=None): + transaction = current_transaction() + if response_id and transaction: + nr_message_ids = getattr(transaction, "_nr_message_ids", {}) + message_id_info = nr_message_ids.pop(response_id, ()) + + if not message_id_info: + warnings.warn("No message ids found for %s" % response_id) + return [] + + conversation_id, request_id, ids = message_id_info + + return [{"conversation_id": conversation_id, "request_id": request_id, "message_id": _id} for _id in ids] + warnings.warn("No message ids found. get_ai_message_ids must be called within the scope of a transaction.") + return [] diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index 336b35ecc9..bff329e554 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -216,7 +216,7 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs): if choices: message_list.extend([choices[0].message]) - create_chat_completion_message_event( + message_ids = create_chat_completion_message_event( transaction, settings.app_name, message_list, @@ -229,6 +229,11 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs): conversation_id, ) + # Cache message ids on transaction for retrieval after open ai call completion. + if not hasattr(transaction, "_nr_message_ids"): + transaction._nr_message_ids = {} + transaction._nr_message_ids[response_id] = message_ids + return response @@ -260,9 +265,12 @@ def create_chat_completion_message_event( request_id, conversation_id, ): + message_ids = [] for index, message in enumerate(message_list): + message_id = "%s-%s" % (response_id, index) + message_ids.append(message_id) chat_completion_message_dict = { - "id": "%s-%s" % (response_id, index), + "id": message_id, "appName": app_name, "conversation_id": conversation_id, "request_id": request_id, @@ -277,6 +285,7 @@ def create_chat_completion_message_event( "vendor": "openAI", } transaction.record_ml_event("LlmChatCompletionMessage", chat_completion_message_dict) + return (conversation_id, request_id, message_ids) async def wrap_embedding_acreate(wrapped, instance, args, kwargs): @@ -456,7 +465,7 @@ async def wrap_chat_completion_acreate(wrapped, instance, args, kwargs): if choices: message_list.extend([choices[0].message]) - create_chat_completion_message_event( + message_ids = create_chat_completion_message_event( transaction, settings.app_name, message_list, @@ -469,6 +478,11 @@ async def wrap_chat_completion_acreate(wrapped, instance, args, kwargs): conversation_id, ) + # Cache message ids on transaction for retrieval after open ai call completion. + if not hasattr(transaction, "_nr_message_ids"): + transaction._nr_message_ids = {} + transaction._nr_message_ids[response_id] = message_ids + return response diff --git a/tests/mlmodel_openai/_mock_external_openai_server.py b/tests/mlmodel_openai/_mock_external_openai_server.py index 3cd2ec9521..44cfb5d0de 100644 --- a/tests/mlmodel_openai/_mock_external_openai_server.py +++ b/tests/mlmodel_openai/_mock_external_openai_server.py @@ -129,6 +129,40 @@ "usage": {"completion_tokens": 11, "prompt_tokens": 53, "total_tokens": 64}, }, ), + "You are a mathematician.": ( + { + "Content-Type": "application/json", + "openai-model": "gpt-3.5-turbo-0613", + "openai-organization": "new-relic-nkmd8b", + "openai-processing-ms": "1469", + "openai-version": "2020-10-01", + "x-ratelimit-limit-requests": "200", + "x-ratelimit-limit-tokens": "40000", + "x-ratelimit-remaining-requests": "199", + "x-ratelimit-remaining-tokens": "39940", + "x-ratelimit-reset-requests": "7m12s", + "x-ratelimit-reset-tokens": "90ms", + "x-request-id": "49dbbffbd3c3f4612aa48def69059aad", + }, + 200, + { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "1 plus 2 is 3.", + "role": "assistant", + }, + } + ], + "created": 1696888865, + "id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat", + "model": "gpt-3.5-turbo-0613", + "object": "chat.completion", + "usage": {"completion_tokens": 11, "prompt_tokens": 53, "total_tokens": 64}, + }, + ), } diff --git a/tests/mlmodel_openai/test_get_ai_message_ids.py b/tests/mlmodel_openai/test_get_ai_message_ids.py new file mode 100644 index 0000000000..1a5c29f878 --- /dev/null +++ b/tests/mlmodel_openai/test_get_ai_message_ids.py @@ -0,0 +1,211 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import openai +from testing_support.fixtures import reset_core_stats_engine + +from newrelic.api.background_task import background_task +from newrelic.api.ml_model import get_ai_message_ids +from newrelic.api.transaction import add_custom_attribute, current_transaction + +_test_openai_chat_completion_messages_1 = ( + {"role": "system", "content": "You are a scientist."}, + {"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"}, +) +_test_openai_chat_completion_messages_2 = ( + {"role": "system", "content": "You are a mathematician."}, + {"role": "user", "content": "What is 1 plus 2?"}, +) +expected_message_ids_1 = [ + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-0", + }, + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-1", + }, + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-2", + }, +] + +expected_message_ids_1_no_conversation_id = [ + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-0", + }, + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-1", + }, + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059ccd", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-2", + }, +] +expected_message_ids_2 = [ + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-0", + }, + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-1", + }, + { + "conversation_id": "my-awesome-id", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-2", + }, +] +expected_message_ids_2_no_conversation_id = [ + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-0", + }, + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-1", + }, + { + "conversation_id": "", + "request_id": "49dbbffbd3c3f4612aa48def69059aad", + "message_id": "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat-2", + }, +] + + +@reset_core_stats_engine() +@background_task() +def test_get_ai_message_ids_when_nr_message_ids_not_set(): + message_ids = get_ai_message_ids("request-id-1") + assert message_ids == [] + + +@reset_core_stats_engine() +def test_get_ai_message_ids_outside_transaction(): + message_ids = get_ai_message_ids("request-id-1") + assert message_ids == [] + + +@reset_core_stats_engine() +@background_task() +def test_get_ai_message_ids_mulitple_async(loop, set_trace_info): + set_trace_info() + add_custom_attribute("conversation_id", "my-awesome-id") + + async def _run(): + res1 = await openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100 + ) + res2 = await openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100 + ) + return [res1, res2] + + results = loop.run_until_complete(_run()) + + message_ids = [m for m in get_ai_message_ids(results[0].id)] + assert message_ids == expected_message_ids_1 + + message_ids = [m for m in get_ai_message_ids(results[1].id)] + assert message_ids == expected_message_ids_2 + + # Make sure we aren't causing a memory leak. + transaction = current_transaction() + assert not transaction._nr_message_ids + + +@reset_core_stats_engine() +@background_task() +def test_get_ai_message_ids_mulitple_async_no_conversation_id(loop, set_trace_info): + set_trace_info() + + async def _run(): + res1 = await openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100 + ) + res2 = await openai.ChatCompletion.acreate( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100 + ) + return [res1, res2] + + results = loop.run_until_complete(_run()) + + message_ids = [m for m in get_ai_message_ids(results[0].id)] + assert message_ids == expected_message_ids_1_no_conversation_id + + message_ids = [m for m in get_ai_message_ids(results[1].id)] + assert message_ids == expected_message_ids_2_no_conversation_id + + # Make sure we aren't causing a memory leak. + transaction = current_transaction() + assert not transaction._nr_message_ids + + +@reset_core_stats_engine() +@background_task() +def test_get_ai_message_ids_mulitple_sync(set_trace_info): + set_trace_info() + add_custom_attribute("conversation_id", "my-awesome-id") + + results = openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100 + ) + message_ids = [m for m in get_ai_message_ids(results.id)] + assert message_ids == expected_message_ids_1 + + results = openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100 + ) + message_ids = [m for m in get_ai_message_ids(results.id)] + assert message_ids == expected_message_ids_2 + + # Make sure we aren't causing a memory leak. + transaction = current_transaction() + assert not transaction._nr_message_ids + + +@reset_core_stats_engine() +@background_task() +def test_get_ai_message_ids_mulitple_sync_no_conversation_id(set_trace_info): + set_trace_info() + + results = openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100 + ) + message_ids = [m for m in get_ai_message_ids(results.id)] + assert message_ids == expected_message_ids_1_no_conversation_id + + results = openai.ChatCompletion.create( + model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_2, temperature=0.7, max_tokens=100 + ) + message_ids = [m for m in get_ai_message_ids(results.id)] + assert message_ids == expected_message_ids_2_no_conversation_id + + # Make sure we aren't causing a memory leak. + transaction = current_transaction() + assert not transaction._nr_message_ids