Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions newrelic/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
38 changes: 38 additions & 0 deletions newrelic/hooks/mlmodel_autogen.py
Original file line number Diff line number Diff line change
@@ -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)
30 changes: 30 additions & 0 deletions tests/mlmodel_autogen/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
# 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
)
67 changes: 67 additions & 0 deletions tests/mlmodel_autogen/test_mcp_tool_adapter.py
Original file line number Diff line number Diff line change
@@ -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.

from unittest.mock import AsyncMock

import pytest
from autogen_ext.tools.mcp import SseMcpToolAdapter, SseServerParams
from mcp import ClientSession, Tool
from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics

from newrelic.api.background_task import background_task

# 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())
4 changes: 4 additions & 0 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down