diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 561d9011f8..ca63991af6 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -158,7 +158,6 @@ def extract_bedrock_titan_text_model(request_body, response_body=None): input_message_list = [{"role": "user", "content": request_body.get("inputText", "")}] - chat_completion_summary_dict = { "request.max_tokens": request_config.get("maxTokenCount", ""), "request.temperature": request_config.get("temperature", ""), @@ -170,7 +169,9 @@ def extract_bedrock_titan_text_model(request_body, response_body=None): completion_tokens = sum(result["tokenCount"] for result in response_body.get("results", [])) total_tokens = input_tokens + completion_tokens - output_message_list = [{"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", [])] + output_message_list = [ + {"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", []) + ] chat_completion_summary_dict.update( { @@ -218,7 +219,9 @@ def extract_bedrock_ai21_j2_model(request_body, response_body=None): } if response_body: - output_message_list =[{"role": "assistant", "content": result["data"]["text"]} for result in response_body.get("completions", [])] + output_message_list = [ + {"role": "assistant", "content": result["data"]["text"]} for result in response_body.get("completions", []) + ] chat_completion_summary_dict.update( { @@ -275,7 +278,9 @@ def extract_bedrock_cohere_model(request_body, response_body=None): } if response_body: - output_message_list = [{"role": "assistant", "content": result["text"]} for result in response_body.get("generations", [])] + output_message_list = [ + {"role": "assistant", "content": result["text"]} for result in response_body.get("generations", []) + ] chat_completion_summary_dict.update( { "response.choices.finish_reason": response_body["generations"][0]["finish_reason"], @@ -377,13 +382,31 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs): if operation == "embedding": # Only available embedding models handle_embedding_event( - instance, transaction, extractor, model, None, None, request_body, - ft.duration, True, trace_id, span_id + instance, + transaction, + extractor, + model, + None, + None, + request_body, + ft.duration, + True, + trace_id, + span_id ) else: handle_chat_completion_event( - instance, transaction, extractor, model, None, None, request_body, - ft.duration, True, trace_id, span_id + instance, + transaction, + extractor, + model, + None, + None, + request_body, + ft.duration, + True, + trace_id, + span_id ) finally: @@ -430,7 +453,17 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs): def handle_embedding_event( - client, transaction, extractor, model, response_body, response_headers, request_body, duration, is_error, trace_id, span_id + client, + transaction, + extractor, + model, + response_body, + response_headers, + request_body, + duration, + is_error, + trace_id, + span_id ): embedding_id = str(uuid.uuid4()) @@ -465,10 +498,20 @@ def handle_embedding_event( def handle_chat_completion_event( - client, transaction, extractor, model, response_body, response_headers, request_body, duration, is_error, trace_id, span_id + client, + transaction, + extractor, + model, + response_body, + response_headers, + request_body, + duration, + is_error, + trace_id, + span_id ): custom_attrs_dict = transaction._custom_params - conversation_id = custom_attrs_dict.get("conversation_id", "") + conversation_id = custom_attrs_dict.get("llm.conversation_id", "") chat_completion_id = str(uuid.uuid4()) diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index babfaf8bab..8534502289 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -193,7 +193,7 @@ def wrap_chat_completion_sync(wrapped, instance, args, kwargs): # Get conversation ID off of the transaction custom_attrs_dict = transaction._custom_params - conversation_id = custom_attrs_dict.get("conversation_id", "") + conversation_id = custom_attrs_dict.get("llm.conversation_id", "") settings = transaction.settings if transaction.settings is not None else global_settings() app_name = settings.app_name @@ -650,7 +650,7 @@ async def wrap_chat_completion_async(wrapped, instance, args, kwargs): # Get conversation ID off of the transaction custom_attrs_dict = transaction._custom_params - conversation_id = custom_attrs_dict.get("conversation_id", "") + conversation_id = custom_attrs_dict.get("llm.conversation_id", "") settings = transaction.settings if transaction.settings is not None else global_settings() app_name = settings.app_name diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index efcc7cec05..2c4925a43b 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -23,7 +23,6 @@ chat_completion_expected_events, chat_completion_invalid_access_key_error_events, chat_completion_payload_templates, - chat_completion_invalid_access_key_error_events, ) from conftest import BOTOCORE_VERSION from testing_support.fixtures import ( @@ -128,7 +127,7 @@ def test_bedrock_chat_completion_in_txn_with_convo_id(set_trace_info, exercise_m @background_task(name="test_bedrock_chat_completion_in_txn_with_convo_id") def _test(): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) _test() @@ -160,7 +159,7 @@ def _test(): @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_bedrock_chat_completion_outside_txn(set_trace_info, exercise_model): - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") exercise_model(prompt=_test_bedrock_chat_completion_prompt, temperature=0.7, max_tokens=100) @@ -237,7 +236,7 @@ def test_bedrock_chat_completion_error_invalid_model(bedrock_server, set_trace_i @background_task(name="test_bedrock_chat_completion_error_invalid_model") def _test(): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") with pytest.raises(_client_error): bedrock_server.invoke_model( body=b"{}", @@ -283,7 +282,7 @@ def _test(): with pytest.raises(_client_error): # not sure where this exception actually comes from set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") exercise_model(prompt="Invalid Token", temperature=0.7, max_tokens=100) _test() diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index e141e45e53..76017a22a8 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -146,7 +146,7 @@ @background_task() def test_openai_chat_completion_sync_in_txn_with_convo_id(set_trace_info): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") openai.ChatCompletion.create( model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) @@ -272,7 +272,7 @@ def test_openai_chat_completion_sync_in_txn_no_convo_id(set_trace_info): @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_openai_chat_completion_sync_outside_txn(): - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") openai.ChatCompletion.create( model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) @@ -335,7 +335,7 @@ def test_openai_chat_completion_async_conversation_id_unset(loop, set_trace_info @background_task() def test_openai_chat_completion_async_conversation_id_set(loop, set_trace_info): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( openai.ChatCompletion.acreate( diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index fe298c02bb..a8d3bdc512 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -131,7 +131,7 @@ def test_chat_completion_invalid_request_error_no_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") openai.ChatCompletion.create( # no model provided, messages=_test_openai_chat_completion_messages, @@ -215,7 +215,7 @@ def test_chat_completion_invalid_request_error_no_model(set_trace_info): def test_chat_completion_invalid_request_error_invalid_model(set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") openai.ChatCompletion.create( model="does-not-exist", messages=({"role": "user", "content": "Model does not exist."},), @@ -315,7 +315,7 @@ def test_chat_completion_invalid_request_error_invalid_model(set_trace_info): def test_chat_completion_authentication_error(monkeypatch, set_trace_info): with pytest.raises(openai.error.AuthenticationError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") monkeypatch.setattr(openai, "api_key", None) # openai.api_key = None openai.ChatCompletion.create( model="gpt-3.5-turbo", @@ -439,7 +439,7 @@ def test_chat_completion_wrong_api_key_error(monkeypatch, set_trace_info): def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( openai.ChatCompletion.acreate( # no model provided, @@ -481,7 +481,7 @@ def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_in def test_chat_completion_invalid_request_error_invalid_model_async(loop, set_trace_info): with pytest.raises(openai.InvalidRequestError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( openai.ChatCompletion.acreate( model="does-not-exist", @@ -520,7 +520,7 @@ def test_chat_completion_invalid_request_error_invalid_model_async(loop, set_tra def test_chat_completion_authentication_error_async(loop, monkeypatch, set_trace_info): with pytest.raises(openai.error.AuthenticationError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") monkeypatch.setattr(openai, "api_key", None) # openai.api_key = None loop.run_until_complete( openai.ChatCompletion.acreate( diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index 70dc58f998..670689c929 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -127,7 +127,7 @@ def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_openai_client): with pytest.raises(TypeError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") sync_openai_client.chat.completions.create( messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) @@ -160,7 +160,7 @@ def test_chat_completion_invalid_request_error_no_model(set_trace_info, sync_ope def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(TypeError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( async_openai_client.chat.completions.create( messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 @@ -242,7 +242,7 @@ def test_chat_completion_invalid_request_error_no_model_async(loop, set_trace_in def test_chat_completion_invalid_request_error_invalid_model(set_trace_info, sync_openai_client): with pytest.raises(openai.NotFoundError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") sync_openai_client.chat.completions.create( model="does-not-exist", messages=({"role": "user", "content": "Model does not exist."},), @@ -281,7 +281,7 @@ def test_chat_completion_invalid_request_error_invalid_model(set_trace_info, syn def test_chat_completion_invalid_request_error_invalid_model_async(loop, set_trace_info, async_openai_client): with pytest.raises(openai.NotFoundError): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( async_openai_client.chat.completions.create( model="does-not-exist", diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index 4df977a6c2..b1b35826c9 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -146,7 +146,7 @@ @background_task() def test_openai_chat_completion_sync_in_txn_with_convo_id(set_trace_info, sync_openai_client): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") sync_openai_client.chat.completions.create( model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) @@ -272,7 +272,7 @@ def test_openai_chat_completion_sync_in_txn_no_convo_id(set_trace_info, sync_ope @reset_core_stats_engine() @validate_custom_event_count(count=0) def test_openai_chat_completion_sync_outside_txn(sync_openai_client): - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") sync_openai_client.chat.completions.create( model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) @@ -335,7 +335,7 @@ def test_openai_chat_completion_async_conversation_id_unset(loop, set_trace_info @background_task() def test_openai_chat_completion_async_conversation_id_set(loop, set_trace_info, async_openai_client): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") loop.run_until_complete( async_openai_client.chat.completions.create( diff --git a/tests/mlmodel_openai/test_get_llm_message_ids.py b/tests/mlmodel_openai/test_get_llm_message_ids.py index af073f7300..8489f4f3d3 100644 --- a/tests/mlmodel_openai/test_get_llm_message_ids.py +++ b/tests/mlmodel_openai/test_get_llm_message_ids.py @@ -13,10 +13,14 @@ # limitations under the License. import openai +from testing_support.fixtures import ( + reset_core_stats_engine, + validate_custom_event_count, +) + from newrelic.api.background_task import background_task from newrelic.api.ml_model import get_llm_message_ids, record_llm_feedback_event from newrelic.api.transaction import add_custom_attribute, current_transaction -from testing_support.fixtures import reset_core_stats_engine, validate_custom_event_count _test_openai_chat_completion_messages_1 = ( {"role": "system", "content": "You are a scientist."}, @@ -114,7 +118,7 @@ def test_get_llm_message_ids_outside_transaction(): @background_task() def test_get_llm_message_ids_mulitple_async(loop, set_trace_info): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") async def _run(): res1 = await openai.ChatCompletion.acreate( @@ -172,7 +176,7 @@ async def _run(): @background_task() def test_get_llm_message_ids_mulitple_sync(set_trace_info): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.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 diff --git a/tests/mlmodel_openai/test_get_llm_message_ids_v1.py b/tests/mlmodel_openai/test_get_llm_message_ids_v1.py index f85a26c2a9..094ddcd5a7 100644 --- a/tests/mlmodel_openai/test_get_llm_message_ids_v1.py +++ b/tests/mlmodel_openai/test_get_llm_message_ids_v1.py @@ -116,7 +116,7 @@ def test_get_llm_message_ids_outside_transaction(): @background_task() def test_get_llm_message_ids_mulitple_async(loop, set_trace_info, async_openai_client): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") async def _run(): res1 = await async_openai_client.chat.completions.create( @@ -174,7 +174,7 @@ async def _run(): @background_task() def test_get_llm_message_ids_mulitple_sync(set_trace_info, sync_openai_client): set_trace_info() - add_custom_attribute("conversation_id", "my-awesome-id") + add_custom_attribute("llm.conversation_id", "my-awesome-id") results = sync_openai_client.chat.completions.create( model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages_1, temperature=0.7, max_tokens=100