From 2de823f8026578e18dc61e4edec803a3e0685949 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Fri, 15 Dec 2023 15:48:58 -0800 Subject: [PATCH] Mark instrumentation points for SDK --- newrelic/hooks/external_botocore.py | 78 +++++++++++++++---- newrelic/hooks/mlmodel_openai.py | 32 +++++--- .../test_bedrock_chat_completion.py | 5 +- .../test_bedrock_embeddings.py | 8 +- tests/mlmodel_openai/test_chat_completion.py | 6 ++ tests/mlmodel_openai/test_embeddings.py | 6 ++ 6 files changed, 108 insertions(+), 27 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 12bdfcafe2..973fa36ea3 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -144,7 +144,7 @@ def create_chat_completion_message_event( "response.model": request_model, "vendor": "bedrock", "ingest_source": "Python", - "is_response": True + "is_response": True, } transaction.record_custom_event("LlmChatCompletionMessage", chat_completion_message_dict) @@ -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( { @@ -243,7 +246,7 @@ def extract_bedrock_claude_model(request_body, response_body=None): chat_completion_summary_dict = { "request.max_tokens": request_body.get("max_tokens_to_sample", ""), "request.temperature": request_body.get("temperature", ""), - "response.number_of_messages": len(input_message_list) + "response.number_of_messages": len(input_message_list), } if response_body: @@ -271,11 +274,13 @@ def extract_bedrock_cohere_model(request_body, response_body=None): chat_completion_summary_dict = { "request.max_tokens": request_body.get("max_tokens", ""), "request.temperature": request_body.get("temperature", ""), - "response.number_of_messages": len(input_message_list) + "response.number_of_messages": len(input_message_list), } 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"], @@ -363,7 +368,7 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs): notice_error_attributes = { "http.statusCode": error_attributes["http.statusCode"], "error.message": error_attributes["error.message"], - "error.code": error_attributes["error.code"] + "error.code": error_attributes["error.code"], } if is_embedding: @@ -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,7 +498,17 @@ 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", "") @@ -549,6 +592,12 @@ def _nr_clientcreator__create_api_method_(wrapped, instance, args, kwargs): return tracer(wrapped) +def _nr_clientcreator__create_methods(wrapped, instance, args, kwargs): + class_attributes = wrapped(*args, **kwargs) + class_attributes["_nr_wrapped"] = True + return class_attributes + + def _bind_make_request_params(operation_model, request_dict, *args, **kwargs): return operation_model, request_dict @@ -579,3 +628,4 @@ def instrument_botocore_endpoint(module): def instrument_botocore_client(module): wrap_function_wrapper(module, "ClientCreator._create_api_method", _nr_clientcreator__create_api_method_) + wrap_function_wrapper(module, "ClientCreator._create_methods", _nr_clientcreator__create_methods) diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py index a653b7ca69..34a9c75d18 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -834,21 +834,33 @@ def wrap_base_client_process_response(wrapped, instance, args, kwargs): def instrument_openai_util(module): - wrap_function_wrapper(module, "convert_to_openai_object", wrap_convert_to_openai_object) + if hasattr(module, "convert_to_openai_object"): + wrap_function_wrapper(module, "convert_to_openai_object", wrap_convert_to_openai_object) + # This is to mark where we instrument so the SDK knows not to instrument them + # again. + setattr(module.convert_to_openai_object, "_nr_wrapped", True) def instrument_openai_api_resources_embedding(module): - if hasattr(module.Embedding, "create"): - wrap_function_wrapper(module, "Embedding.create", wrap_embedding_sync) - if hasattr(module.Embedding, "acreate"): - wrap_function_wrapper(module, "Embedding.acreate", wrap_embedding_async) + if hasattr(module, "Embedding"): + if hasattr(module.Embedding, "create"): + wrap_function_wrapper(module, "Embedding.create", wrap_embedding_sync) + if hasattr(module.Embedding, "acreate"): + wrap_function_wrapper(module, "Embedding.acreate", wrap_embedding_async) + # This is to mark where we instrument so the SDK knows not to instrument them + # again. + setattr(module.Embedding, "_nr_wrapped", True) def instrument_openai_api_resources_chat_completion(module): - if hasattr(module.ChatCompletion, "create"): - wrap_function_wrapper(module, "ChatCompletion.create", wrap_chat_completion_sync) - if hasattr(module.ChatCompletion, "acreate"): - wrap_function_wrapper(module, "ChatCompletion.acreate", wrap_chat_completion_async) + if hasattr(module, "ChatCompletion"): + if hasattr(module.ChatCompletion, "create"): + wrap_function_wrapper(module, "ChatCompletion.create", wrap_chat_completion_sync) + if hasattr(module.ChatCompletion, "acreate"): + wrap_function_wrapper(module, "ChatCompletion.acreate", wrap_chat_completion_async) + # This is to mark where we instrument so the SDK knows not to instrument them + # again. + setattr(module.ChatCompletion, "_nr_wrapped", True) def instrument_openai_resources_chat_completions(module): @@ -858,7 +870,6 @@ def instrument_openai_resources_chat_completions(module): wrap_function_wrapper(module, "AsyncCompletions.create", wrap_chat_completion_async) -# OpenAI v1 instrumentation points def instrument_openai_resources_embeddings(module): if hasattr(module, "Embeddings"): if hasattr(module.Embeddings, "create"): @@ -872,3 +883,4 @@ def instrument_openai_resources_embeddings(module): def instrument_openai_base_client(module): if hasattr(module.BaseClient, "_process_response"): wrap_function_wrapper(module, "BaseClient._process_response", wrap_base_client_process_response) + setattr(module.BaseClient._process_response, "_nr_wrapped", True) diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 604771c824..f51fbefda4 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 ( @@ -287,3 +286,7 @@ def _test(): exercise_model(prompt="Invalid Token", temperature=0.7, max_tokens=100) _test() + + +def test_bedrock_chat_completion_functions_marked_as_wrapped_for_sdk_compatibility(bedrock_server): + assert bedrock_server._nr_wrapped diff --git a/tests/external_botocore/test_bedrock_embeddings.py b/tests/external_botocore/test_bedrock_embeddings.py index 7a5740e465..9fc0164714 100644 --- a/tests/external_botocore/test_bedrock_embeddings.py +++ b/tests/external_botocore/test_bedrock_embeddings.py @@ -1,4 +1,4 @@ - # Copyright 2010 New Relic, Inc. +# 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. @@ -19,8 +19,8 @@ import pytest from _test_bedrock_embeddings import ( embedding_expected_client_errors, - embedding_expected_events, embedding_expected_error_events, + embedding_expected_events, embedding_payload_templates, ) from conftest import BOTOCORE_VERSION @@ -172,3 +172,7 @@ def _test(): exercise_model(prompt="Invalid Token", temperature=0.7, max_tokens=100) _test() + + +def test_bedrock_chat_completion_functions_marked_as_wrapped_for_sdk_compatibility(bedrock_server): + assert bedrock_server._nr_wrapped diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index f2c31b2628..cd7750e329 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -371,3 +371,9 @@ def test_openai_chat_completion_async_disabled_custom_event_settings(loop): model="gpt-3.5-turbo", messages=_test_openai_chat_completion_messages, temperature=0.7, max_tokens=100 ) ) + + +def test_openai_chat_completion_functions_marked_as_wrapped_for_sdk_compatibility(): + assert openai.ChatCompletion._nr_wrapped + assert openai.ChatCompletion._nr_wrapped + assert openai.util.convert_to_openai_object diff --git a/tests/mlmodel_openai/test_embeddings.py b/tests/mlmodel_openai/test_embeddings.py index ae2c048fc2..7dc4d6d888 100644 --- a/tests/mlmodel_openai/test_embeddings.py +++ b/tests/mlmodel_openai/test_embeddings.py @@ -148,3 +148,9 @@ def test_openai_embedding_async_disabled_custom_insights_events(loop): loop.run_until_complete( openai.Embedding.acreate(input="This is an embedding test.", model="text-embedding-ada-002") ) + + +def test_openai_embedding_functions_marked_as_wrapped_for_sdk_compatibility(): + assert openai.Embedding._nr_wrapped + assert openai.Embedding._nr_wrapped + assert openai.util.convert_to_openai_object