Skip to content

Commit

Permalink
Mark instrumentation points for SDK (#1009)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
hmstepanek and mergify[bot] committed Feb 13, 2024
1 parent 9ffd241 commit b8c3a38
Show file tree
Hide file tree
Showing 6 changed files with 47 additions and 11 deletions.
7 changes: 7 additions & 0 deletions newrelic/hooks/external_botocore.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,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

Expand Down Expand Up @@ -622,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)
31 changes: 21 additions & 10 deletions newrelic/hooks/mlmodel_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -884,21 +884,33 @@ async def wrap_base_client_process_response_async(wrapped, instance, args, kwarg


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):
Expand All @@ -908,7 +920,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"):
Expand Down
4 changes: 4 additions & 0 deletions tests/external_botocore/test_bedrock_chat_completion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 5 additions & 1 deletion tests/external_botocore/test_bedrock_embeddings.py
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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
5 changes: 5 additions & 0 deletions tests/mlmodel_openai/test_chat_completion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 5 additions & 0 deletions tests/mlmodel_openai/test_embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit b8c3a38

Please sign in to comment.