Skip to content

Commit

Permalink
feat: Warn when agent is not instrumenting
Browse files Browse the repository at this point in the history
  • Loading branch information
RyanKadri committed Dec 21, 2023
1 parent 775e207 commit ede93e1
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 13 deletions.
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ keywords = [

[tool.poetry.dependencies]
python = ">=3.8.1,<4.0.0"
newrelic = { git = "https://github.com/newrelic/newrelic-python-agent.git", rev = "d74a2d7ef0914cb5ae97028b00f65b093e1a9d73" }
newrelic = { git = "https://github.com/newrelic/newrelic-python-agent.git", rev = "7051455f076d2f617719d9bf6ddbca6989fa0afd" }
openai = ">=0.8,<0.30"
tiktoken = "^0.5.1"

Expand Down
26 changes: 23 additions & 3 deletions src/nr_openai_observability/patcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,20 @@
logger = logging.getLogger("nr_openai_observability")


def patched_call(parent_class, method_name, patched_fn, stream_patched_fn=None):
def patched_call(
parent_class,
method_name,
patched_fn,
stream_patched_fn=None,
expect_agent_instrumentation=False,
):
original_fn = getattr(parent_class, method_name)
patched_by_plugin = hasattr(original_fn, "is_patched_by_monitor")
# TODO - This seems to work but it feels like a bit of a hack. Need to coordinate on a better approach with agent team
patched_by_agent = getattr(parent_class, "_nr_wrapped", False)
if expect_agent_instrumentation and not patched_by_agent:
logger.warn(
f"{parent_class.__name__}{method_name} should be instrumented by the python agent but it is not. Falling back to plugin instrumentation"
)

if patched_by_plugin:
return original_fn
Expand All @@ -32,11 +41,22 @@ def _inner_patch(*args, **kwargs):
return _inner_patch


def patched_call_async(parent_class, method_name, patched_fn, stream_patched_fn=None):
def patched_call_async(
parent_class,
method_name,
patched_fn,
stream_patched_fn=None,
expect_agent_instrumentation=False,
):
original_fn = getattr(parent_class, method_name)
patched_by_plugin = hasattr(original_fn, "is_patched_by_monitor")
patched_by_agent = getattr(parent_class, "_nr_wrapped", False)

if expect_agent_instrumentation and not patched_by_agent:
logger.warn(
f"{str(parent_class)}{method_name} should be instrumented by the python agent but it is not. Falling back to plugin instrumentation"
)

if patched_by_plugin:
return original_fn

Expand Down
10 changes: 8 additions & 2 deletions src/nr_openai_observability/patchers/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,10 @@ def patcher_aws_create_api(original_fn, *args, **kwargs):
response,
bedrock_method,
patched_call(
response, bedrock_method, patcher_bedrock_create_completion
response,
bedrock_method,
patcher_bedrock_create_completion,
expect_agent_instrumentation=True,
),
)
else:
Expand Down Expand Up @@ -558,7 +561,10 @@ def perform_patch_bedrock():

try:
botocore.client.ClientCreator.create_client = patched_call(
botocore.client.ClientCreator, "create_client", patcher_aws_create_api
botocore.client.ClientCreator,
"create_client",
patcher_aws_create_api,
expect_agent_instrumentation=False,
)
except AttributeError as error:
logger.debug(
Expand Down
5 changes: 4 additions & 1 deletion src/nr_openai_observability/patchers/langchain.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,10 @@ def perform_patch_langchain_vectorstores():
try:
cls = getattr(langchain.vectorstores, vector_store)
cls.similarity_search = patched_call(
cls, "similarity_search", patcher_similarity_search
cls,
"similarity_search",
patcher_similarity_search,
expect_agent_instrumentation=False,
)
except AttributeError:
pass
28 changes: 22 additions & 6 deletions src/nr_openai_observability/patchers/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,6 @@ def perform_patch_openai():
(major, minor, revision) = openai_version
too_new = major > 1
requires_agent = major == 1
agent_or_plugin = major == 0 and minor >= 26
too_old = major == 0 and minor < MIN_MINOR_VERSION

supported_versions_msg = f"Versions between v{MIN_MAJOR_VERSION}.${MIN_MINOR_VERSION}.0 and <{MAX_MAJOR_VERSION + 1}.0 are supported"
Expand All @@ -338,28 +337,40 @@ def perform_patch_openai():

try:
openai.Embedding.create = patched_call(
openai.Embedding, "create", patcher_create_embedding
openai.Embedding,
"create",
patcher_create_embedding,
expect_agent_instrumentation=True,
)
except AttributeError:
pass

try:
openai.Embedding.acreate = patched_call_async(
openai.Embedding, "acreate", patcher_create_embedding_async
openai.Embedding,
"acreate",
patcher_create_embedding_async,
expect_agent_instrumentation=True,
)
except AttributeError:
pass

try:
openai.Completion.create = patched_call(
openai.Completion, "create", patcher_create_completion
openai.Completion,
"create",
patcher_create_completion,
expect_agent_instrumentation=False,
)
except AttributeError:
pass

try:
openai.Completion.acreate = patched_call_async(
openai.Completion, "acreate", patcher_create_completion_async
openai.Completion,
"acreate",
patcher_create_completion_async,
expect_agent_instrumentation=False,
)
except AttributeError:
pass
Expand All @@ -370,6 +381,7 @@ def perform_patch_openai():
"create",
patcher_create_chat_completion,
patcher_create_chat_completion_stream,
expect_agent_instrumentation=True,
)
except AttributeError:
pass
Expand All @@ -380,13 +392,17 @@ def perform_patch_openai():
"acreate",
patcher_create_chat_completion_async,
patcher_create_chat_completion_stream_async,
expect_agent_instrumentation=True,
)
except AttributeError:
pass

try:
openai.util.convert_to_openai_object = patched_call(
openai.util, "convert_to_openai_object", patcher_convert_to_openai_object
openai.util,
"convert_to_openai_object",
patcher_convert_to_openai_object,
expect_agent_instrumentation=False,
)
except AttributeError:
pass

0 comments on commit ede93e1

Please sign in to comment.