Test for remote tools in langchain

Author: TimPansino

tests/mlmodel_langchain/_test_tools.py

@@ -32,6 +32,23 @@ async def add_exclamation_async(message: str) -> str:
     return f"{message}!"
 
 
+@tool("add_exclamation")
+def remote_add_exclamation_sync(message: str) -> str:
+    """Adds an exclamation mark to the input message."""
+    return f"{message}!"
+
+
+@tool("add_exclamation")
+async def remote_add_exclamation_async(message: str) -> str:
+    """Adds an exclamation mark to the input message."""
+    return f"{message}!"
+
+
+# Simulate a remote tool by setting the __module__ to a known remote tool module
+remote_add_exclamation_sync.func.__module__ = "langchain_mcp_adapters.tools"
+remote_add_exclamation_async.coroutine.__module__ = "langchain_mcp_adapters.tools"
+
+
 @pytest.fixture(params=["sync_tool", "async_tool"])
 def tool_type(request):
     return request.param
@@ -52,3 +69,15 @@ def add_exclamation(tool_type, exercise_agent):
         return add_exclamation_async
     else:
         raise NotImplementedError
+
+
+@pytest.fixture
+def remote_add_exclamation(tool_type, exercise_agent):
+    if tool_type == "sync_tool":
+        return remote_add_exclamation_sync
+    elif tool_type == "async_tool":
+        if exercise_agent._called_method in ("invoke", "stream"):
+            pytest.skip("Async tools cannot be invoked synchronously.")
+        return remote_add_exclamation_async
+    else:
+        raise NotImplementedError

tests/mlmodel_langchain/test_tools.py

@@ -32,7 +32,16 @@
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import transient_function_wrapper
 
-from ._test_tools import add_exclamation, tool_method_name, tool_type
+from ._test_tools import (
+    add_exclamation,
+    add_exclamation_async,
+    add_exclamation_sync,
+    remote_add_exclamation,
+    remote_add_exclamation_async,
+    remote_add_exclamation_sync,
+    tool_method_name,
+    tool_type,
+)
 
 PROMPT = {
     "messages": [
@@ -155,6 +164,38 @@ def _test():
     _test()
 
 
+@dt_enabled
+@reset_core_stats_engine()
+def test_tool_remote_no_subcomponent(
+    exercise_agent, set_trace_info, create_agent_runnable, remote_add_exclamation, tool_method_name
+):
+    @validate_custom_events(tool_recorded_event)
+    @validate_custom_event_count(count=exercise_agent._expected_event_count)
+    @validate_transaction_metrics(
+        "test_tool_remote_no_subcomponent",
+        scoped_metrics=[(f"Llm/tool/LangChain/{tool_method_name}/add_exclamation", 1)],
+        rollup_metrics=[(f"Llm/tool/LangChain/{tool_method_name}/add_exclamation", 1)],
+        background_task=True,
+    )
+    # The agent span still has a subcomponent attribute
+    @validate_span_events(count=1, exact_agents={"subcomponent": '{"type": "APM-AI_AGENT", "name": "my_agent"}'})
+    # Tool span must NOT have a subcomponent attribute
+    @validate_span_events(
+        count=1,
+        exact_intrinsics={"name": f"Llm/tool/LangChain/{tool_method_name}/add_exclamation"},
+        unexpected_agents=["subcomponent"],
+    )
+    @background_task(name="test_tool_remote_no_subcomponent")
+    def _test():
+        set_trace_info()
+        my_agent = create_agent_runnable(
+            tools=[remote_add_exclamation], system_prompt="You are a text manipulation algorithm."
+        )
+        exercise_agent(my_agent, PROMPT)
+
+    _test()
+
+
 @dt_enabled
 @reset_core_stats_engine()
 def test_tool_execution_error(exercise_agent, set_trace_info, create_agent_runnable, add_exclamation, tool_method_name):
@@ -219,3 +260,23 @@ def _test():
             exercise_agent(my_agent, PROMPT)
 
     _test()
+
+
+def test_is_local_tool_helper():
+    from newrelic.hooks.mlmodel_langchain import _is_local_tool
+
+    # Local tool functions
+    assert _is_local_tool(add_exclamation_sync) is True
+    assert _is_local_tool(add_exclamation_async) is True
+
+    # Tools whose __module__ reads as a collection of remote tools
+    assert _is_local_tool(remote_add_exclamation_sync) is False
+    assert _is_local_tool(remote_add_exclamation_async) is False
+
+    # A tool instance that raises an exception when trying to access its underlying callable
+    class BrokenTool:
+        @property
+        def func(self):
+            raise RuntimeError("Oops")
+
+    assert _is_local_tool(BrokenTool()) is True