From 5538fa1bf2d4e17e236df216472a1a831038de8e Mon Sep 17 00:00:00 2001 From: Uma Annamalai Date: Mon, 16 Jun 2025 19:27:51 -0700 Subject: [PATCH 1/2] Add instrumentation for MCP tool adapter. --- newrelic/config.py | 4 ++ newrelic/hooks/mlmodel_autogen.py | 38 +++++++++++ tests/mlmodel_autogen/conftest.py | 31 +++++++++ .../mlmodel_autogen/test_mcp_tool_adapter.py | 67 +++++++++++++++++++ tox.ini | 4 ++ 5 files changed, 144 insertions(+) create mode 100644 newrelic/hooks/mlmodel_autogen.py create mode 100644 tests/mlmodel_autogen/conftest.py create mode 100644 tests/mlmodel_autogen/test_mcp_tool_adapter.py diff --git a/newrelic/config.py b/newrelic/config.py index 97b6475c0..f72cf3c70 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2862,6 +2862,10 @@ def _process_module_builtin_defaults(): _process_module_definition("loguru", "newrelic.hooks.logger_loguru", "instrument_loguru") _process_module_definition("loguru._logger", "newrelic.hooks.logger_loguru", "instrument_loguru_logger") + _process_module_definition( + "autogen_ext.tools.mcp._base", "newrelic.hooks.mlmodel_autogen", "instrument_autogen_ext_tools_mcp__base" + ) + _process_module_definition("mcp.client.session", "newrelic.hooks.adapter_mcp", "instrument_mcp_client_session") _process_module_definition("structlog._base", "newrelic.hooks.logger_structlog", "instrument_structlog__base") diff --git a/newrelic/hooks/mlmodel_autogen.py b/newrelic/hooks/mlmodel_autogen.py new file mode 100644 index 000000000..c7b09bf90 --- /dev/null +++ b/newrelic/hooks/mlmodel_autogen.py @@ -0,0 +1,38 @@ +# 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from newrelic.api.function_trace import FunctionTrace +from newrelic.api.transaction import current_transaction +from newrelic.common.object_names import callable_name +from newrelic.common.object_wrapper import wrap_function_wrapper +from newrelic.common.signature import bind_args + + +async def wrap_from_server_params(wrapped, instance, args, kwargs): + transaction = current_transaction() + if not transaction: + return await wrapped(*args, **kwargs) + + func_name = callable_name(wrapped) + bound_args = bind_args(wrapped, args, kwargs) + tool_name = bound_args.get("tool_name") or "tool" + function_trace_name = f"{func_name}/{tool_name}" + with FunctionTrace(name=function_trace_name, group="Llm", source=wrapped): + return await wrapped(*args, **kwargs) + + +def instrument_autogen_ext_tools_mcp__base(module): + if hasattr(module, "McpToolAdapter"): + if hasattr(module.McpToolAdapter, "from_server_params"): + wrap_function_wrapper(module, "McpToolAdapter.from_server_params", wrap_from_server_params) diff --git a/tests/mlmodel_autogen/conftest.py b/tests/mlmodel_autogen/conftest.py new file mode 100644 index 000000000..406ded94a --- /dev/null +++ b/tests/mlmodel_autogen/conftest.py @@ -0,0 +1,31 @@ +# 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from testing_support.fixture.event_loop import event_loop as loop +from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture + + +_default_settings = { + "package_reporting.enabled": False, + "transaction_tracer.explain_threshold": 0.0, + "transaction_tracer.transaction_threshold": 0.0, + "transaction_tracer.stack_trace_threshold": 0.0, + "debug.log_data_collector_payloads": True, + "debug.record_transaction_failure": True, + "ai_monitoring.enabled": True, +} + +collector_agent_registration = collector_agent_registration_fixture( + app_name="Python Agent Test (mlmodel_autogen)", default_settings=_default_settings +) diff --git a/tests/mlmodel_autogen/test_mcp_tool_adapter.py b/tests/mlmodel_autogen/test_mcp_tool_adapter.py new file mode 100644 index 000000000..36b46f2a7 --- /dev/null +++ b/tests/mlmodel_autogen/test_mcp_tool_adapter.py @@ -0,0 +1,67 @@ +# 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from autogen_ext.tools.mcp import SseMcpToolAdapter, SseServerParams +from mcp import ClientSession, Tool +from unittest.mock import AsyncMock + +from newrelic.api.background_task import background_task +from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics + +# Test setup derived from: https://github.com/microsoft/autogen/blob/main/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +# autogen MIT license: https://github.com/microsoft/autogen/blob/main/LICENSE and +# https://github.com/microsoft/autogen/blob/main/LICENSE-CODE + + +@pytest.fixture +def mock_sse_session(): + session = AsyncMock(spec=ClientSession) + session.initialize = AsyncMock() + session.call_tool = AsyncMock() + session.list_tools = AsyncMock() + return session + + +@pytest.fixture +def add_exclamation(): + return Tool( + name="add_exclamation", + description="A test SSE tool that adds an exclamation mark to a string", + inputSchema={"type": "object", "properties": {"input": {"type": "string"}}, "required": ["input"]}, + ) + + +@validate_transaction_metrics( + "test_mcp_tool_adapter:test_from_server_params_tracing", + scoped_metrics=[("Llm/autogen_ext.tools.mcp._sse:SseMcpToolAdapter.from_server_params/add_exclamation", 1)], + rollup_metrics=[("Llm/autogen_ext.tools.mcp._sse:SseMcpToolAdapter.from_server_params/add_exclamation", 1)], + background_task=True, +) +@background_task() +def test_from_server_params_tracing(loop, mock_sse_session, monkeypatch, add_exclamation): + async def _test(): + params = SseServerParams(url="http://test-url") + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_sse_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", lambda *args, **kwargs: mock_context + ) + + mock_sse_session.list_tools.return_value.tools = [add_exclamation] + + adapter = await SseMcpToolAdapter.from_server_params(params, "add_exclamation") + + loop.run_until_complete(_test()) diff --git a/tox.ini b/tox.ini index 23435effb..a025352a8 100644 --- a/tox.ini +++ b/tox.ini @@ -154,6 +154,7 @@ envlist = python-logger_logging-{py37,py38,py39,py310,py311,py312,py313,pypy310}, python-logger_loguru-{py37,py38,py39,py310,py311,py312,py313,pypy310}-logurulatest, python-logger_structlog-{py37,py38,py39,py310,py311,py312,py313,pypy310}-structloglatest, + python-mlmodel_autogen-{py310,py311,py312,py313,pypy310}, python-mlmodel_gemini-{py39,py310,py311,py312,py313}, python-mlmodel_langchain-{py39,py310,py311,py312}, ;; Package not ready for Python 3.13 (uses an older version of numpy) @@ -399,6 +400,8 @@ deps = framework_tornado: pycurl framework_tornado-tornadolatest: tornado framework_tornado-tornadomaster: https://github.com/tornadoweb/tornado/archive/master.zip + mlmodel_autogen: autogen-ext + mlmodel_autogen: mcp mlmodel_gemini: google-genai mlmodel_openai-openai0: openai[datalib]<1.0 mlmodel_openai-openai107: openai[datalib]<1.8 @@ -550,6 +553,7 @@ changedir = messagebroker_kafkapython: tests/messagebroker_kafkapython messagebroker_kombu: tests/messagebroker_kombu messagebroker_pika: tests/messagebroker_pika + mlmodel_autogen: tests/mlmodel_autogen mlmodel_gemini: tests/mlmodel_gemini mlmodel_langchain: tests/mlmodel_langchain mlmodel_openai: tests/mlmodel_openai From 6c38559c1ca997d13bf37ae79be3453bf72c493b Mon Sep 17 00:00:00 2001 From: umaannamalai <19895951+umaannamalai@users.noreply.github.com> Date: Tue, 17 Jun 2025 02:35:44 +0000 Subject: [PATCH 2/2] [MegaLinter] Apply linters fixes --- tests/mlmodel_autogen/conftest.py | 1 - tests/mlmodel_autogen/test_mcp_tool_adapter.py | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/mlmodel_autogen/conftest.py b/tests/mlmodel_autogen/conftest.py index 406ded94a..9d8a7702c 100644 --- a/tests/mlmodel_autogen/conftest.py +++ b/tests/mlmodel_autogen/conftest.py @@ -15,7 +15,6 @@ from testing_support.fixture.event_loop import event_loop as loop from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture - _default_settings = { "package_reporting.enabled": False, "transaction_tracer.explain_threshold": 0.0, diff --git a/tests/mlmodel_autogen/test_mcp_tool_adapter.py b/tests/mlmodel_autogen/test_mcp_tool_adapter.py index 36b46f2a7..af8684b20 100644 --- a/tests/mlmodel_autogen/test_mcp_tool_adapter.py +++ b/tests/mlmodel_autogen/test_mcp_tool_adapter.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest +from unittest.mock import AsyncMock +import pytest from autogen_ext.tools.mcp import SseMcpToolAdapter, SseServerParams from mcp import ClientSession, Tool -from unittest.mock import AsyncMock +from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics from newrelic.api.background_task import background_task -from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics # Test setup derived from: https://github.com/microsoft/autogen/blob/main/python/packages/autogen-ext/tests/tools/test_mcp_tools.py # autogen MIT license: https://github.com/microsoft/autogen/blob/main/LICENSE and