From d74a2d7ef0914cb5ae97028b00f65b093e1a9d73 Mon Sep 17 00:00:00 2001 From: Hannah Stepanek Date: Mon, 18 Dec 2023 10:20:12 -0800 Subject: [PATCH] Mark instrumentation points for SDK (#1009) * Mark instrumentation points for SDK * Remove duplicated assertion * Fixup: assert attribute not function --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- newrelic/hooks/external_botocore.py | 7 +++++ newrelic/hooks/mlmodel_openai.py | 31 +++++++++++++------ .../test_bedrock_chat_completion.py | 4 +++ .../test_bedrock_embeddings.py | 6 +++- tests/mlmodel_openai/test_chat_completion.py | 5 +++ tests/mlmodel_openai/test_embeddings.py | 5 +++ 6 files changed, 47 insertions(+), 11 deletions(-) diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 12bdfcafe..561d9011f 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -549,6 +549,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 +585,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 7ea277e76..babfaf8ba 100644 --- a/newrelic/hooks/mlmodel_openai.py +++ b/newrelic/hooks/mlmodel_openai.py @@ -870,21 +870,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): @@ -894,7 +906,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"): diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index 604771c82..efcc7cec0 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -287,3 +287,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 7a5740e46..cc442fc15 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. @@ -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 f2c31b262..e141e45e5 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -371,3 +371,8 @@ 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.util.convert_to_openai_object._nr_wrapped diff --git a/tests/mlmodel_openai/test_embeddings.py b/tests/mlmodel_openai/test_embeddings.py index ae2c048fc..65ac33e87 100644 --- a/tests/mlmodel_openai/test_embeddings.py +++ b/tests/mlmodel_openai/test_embeddings.py @@ -148,3 +148,8 @@ 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.util.convert_to_openai_object._nr_wrapped