Mark instrumentation points for SDK

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)

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)

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

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

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