Mark instrumentation points for SDK

Author: hmstepanek

newrelic/hooks/external_botocore.py

@@ -243,7 +243,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:
@@ -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)

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"):

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

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

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

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