Handle name collisions between local and remote tools (#52)

Author: Ickerday

.basedpyright/baseline.json

@@ -201,30 +201,6 @@
             }
         ],
         "./splunklib/ai/tools.py": [
-            {
-                "code": "reportUnusedImport",
-                "range": {
-                    "startColumn": 7,
-                    "endColumn": 22,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 43,
-                    "endColumn": 75,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnannotatedClassAttribute",
-                "range": {
-                    "startColumn": 13,
-                    "endColumn": 27,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownVariableType",
                 "range": {
@@ -30169,22 +30145,6 @@
             }
         ],
         "./tests/integration/ai/test_agent_mcp_tools.py": [
-            {
-                "code": "reportUnusedImport",
-                "range": {
-                    "startColumn": 21,
-                    "endColumn": 30,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 4,
-                    "endColumn": 24,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownArgumentType",
                 "range": {

examples/ai_custom_search_app/metadata/local.meta

@@ -86,14 +86,15 @@ reportUnusedCallResult = false
 [tool.ruff.lint]
 fixable = ["ALL"]
 select = [
-  "ANN", # flake8 type annotations
+  "ANN", # flake-8-annotations
   "C4",  # comprehensions
   "DOC", # pydocstyle
   "E",   # pycodestyle
   "F",   # pyflakes
   "I",   # isort
-  "UP",  # pyupgrade
+  "PT",  # flake-8-pytest-rules
   "RUF", # ruff-specific rules
+  "UP",  # pyupgrade
 ]
 
 [tool.ruff.lint.isort]

pyproject.toml

@@ -29,6 +29,7 @@
 from splunklib.ai.tool_filtering import ToolFilters, filter_tools
 from splunklib.ai.tools import (
     Tool,
+    ToolType,
     build_local_tools_path,
     connect_local_mcp,
     connect_remote_mcp,
@@ -171,7 +172,11 @@ async def _start_agent(self) -> AsyncGenerator[Self]:
                     )
                     self.logger.debug("Loading local tools")
                     local_tools = await load_mcp_tools(
-                        local_session, "local", app_id, self.trace_id, self._service
+                        local_session,
+                        ToolType.LOCAL,
+                        app_id,
+                        self.trace_id,
+                        self._service,
                     )
                     self.logger.debug(f"Local tools loaded; {local_tools=}")
                     tools.extend(local_tools)
@@ -188,7 +193,7 @@ async def _start_agent(self) -> AsyncGenerator[Self]:
                     self.logger.debug("Loading remote tools - MCP Server available")
                     remote_tools = await load_mcp_tools(
                         remote_session,
-                        "remote",
+                        ToolType.REMOTE,
                         app_id,
                         self.trace_id,
                         self._service,

splunklib/ai/agent.py

@@ -17,7 +17,7 @@
 
 def get_backend() -> Backend:
     """Get a backend instance."""
-
+    # Lazy import to avoid circular dependency hell between LangChain and SDK
     from splunklib.ai.engines.langchain import langchain_backend_factory
 
     # NOTE: For now we're just using the langchain backend implementation

splunklib/ai/core/backend_registry.py

@@ -87,7 +87,7 @@
     tool_middleware,
 )
 from splunklib.ai.model import OpenAIModel, PredefinedModel
-from splunklib.ai.tools import Tool, ToolException
+from splunklib.ai.tools import Tool, ToolException, ToolType
 
 # Represents a prefix reserved only for internal use.
 # No user-visible tool or subagent name can be prefixed with it.
@@ -102,6 +102,10 @@
 # backward compatibility measure - we're free to use any prefixed tool name.
 CONFLICTING_TOOL_PREFIX = f"{RESERVED_LC_TOOL_PREFIX}tool-"
 
+# Prepended to a local tool name when passed to LangChain to both avoid name conflicts
+# and to allow recovering tool type during LC -> SDK conversion
+LOCAL_TOOL_PREFIX = f"{RESERVED_LC_TOOL_PREFIX}local-"
+
 AGENT_AS_TOOLS_PROMPT = f"""
 You are provided with Agents.
 Agents are more advanced TOOLS, which start with "{AGENT_PREFIX}" prefix.
@@ -242,16 +246,25 @@ async def invoke_agent(req: AgentRequest) -> AgentResponse[Any | None]:
             )
 
 
+def _prepare_langchain_tools(agent_tools: Sequence[Tool]) -> list[BaseTool]:
+    """We prefix every local tool name."""
+    tools = list[BaseTool]()
+    for a_tool in agent_tools:
+        tools.append(_create_langchain_tool(a_tool))
+
+    return tools
+
+
 @final
 class LangChainBackend(Backend):
     @override
     async def create_agent(
         self,
         agent: BaseAgent[OutputT],
     ) -> AgentImpl[OutputT]:
-        system_prompt = agent.system_prompt
-        tools = [_create_langchain_tool(t) for t in agent.tools]
+        tools = _prepare_langchain_tools(agent.tools)
 
+        system_prompt = agent.system_prompt
         if agent.agents:
             seen_names: set[str] = set()
             for subagent in agent.agents:
@@ -466,7 +479,8 @@ def _convert_tool_request_to_lc(
 
 
 def _convert_subagent_request_to_lc(
-    request: SubagentRequest, original_request: LC_ToolCallRequest
+    request: SubagentRequest,
+    original_request: LC_ToolCallRequest,
 ) -> LC_ToolCallRequest:
     return original_request.override(
         tool_call=_map_tool_call_to_langchain(request.call),
@@ -475,7 +489,8 @@ def _convert_subagent_request_to_lc(
 
 
 def _convert_model_request_to_lc(
-    request: ModelRequest, original_request: LC_ModelRequest
+    request: ModelRequest,
+    original_request: LC_ModelRequest,
 ) -> LC_ModelRequest:
     return original_request.override(
         system_message=LC_SystemMessage(content=request.system_message),
@@ -504,7 +519,7 @@ def _convert_tool_message_to_lc(
         case SubagentMessage():
             name = _normalize_agent_name(message.name)
         case ToolMessage():
-            name = _normalize_tool_name(message.name)
+            name = _normalize_tool_name(message.name, message.type)
 
     return LC_ToolMessage(
         name=name,
@@ -515,11 +530,10 @@ def _convert_tool_message_to_lc(
 
 
 def _convert_tool_response_to_lc(
-    response: ToolResponse,
-    call: ToolCall,
+    response: ToolResponse, call: ToolCall
 ) -> LC_ToolMessage:
     return LC_ToolMessage(
-        name=_normalize_tool_name(call.name),
+        name=_normalize_tool_name(call.name, call.type),
         content=response.content,
         tool_call_id=call.id,
         status=response.status,
@@ -554,11 +568,18 @@ def _convert_tool_message_from_lc(
             assert message.name is not None, (
                 "LangChain responded with a nameless tool call"
             )
+
+            tool_type: ToolType = (
+                ToolType.LOCAL
+                if message.name.startswith(LOCAL_TOOL_PREFIX)
+                else ToolType.REMOTE
+            )
             return ToolMessage(
                 name=_denormalize_tool_name(message.name),
                 content=message.content.__str__(),
                 call_id=message.tool_call_id,
                 status=message.status,
+                type=tool_type,
             )
         case LC_Command():
             # NOTE: for now the command is not implemented
@@ -668,7 +689,7 @@ async def _tool_call(**kwargs: dict[str, Any]) -> dict[str, Any] | list[str]:
         except ToolException as e:
             raise LC_ToolException(*e.args) from e
         except LC_ToolException:
-            assert False, (
+            assert False, (  # noqa: PT015
                 "ToolException from LangChain should not be raised in tool.func"
             )
 
@@ -687,7 +708,7 @@ async def _tool_call(**kwargs: dict[str, Any]) -> dict[str, Any] | list[str]:
         return result.content
 
     return StructuredTool(
-        name=_normalize_tool_name(tool.name),
+        name=_normalize_tool_name(tool.name, tool.type),
         description=tool.description,
         args_schema=tool.input_schema,
         coroutine=_tool_call,
@@ -709,14 +730,24 @@ def _denormalize_agent_name(name: str) -> str:
     return name.removeprefix(AGENT_PREFIX)
 
 
-def _normalize_tool_name(name: str) -> str:
+def _normalize_tool_name(name: str, tool_type: ToolType) -> str:
+    if tool_type == ToolType.LOCAL:
+        return LOCAL_TOOL_PREFIX + name
+
     if name.startswith(RESERVED_LC_TOOL_PREFIX):
-        return f"{CONFLICTING_TOOL_PREFIX}{name}"
+        # Tool name contains our reserved prefix, see comment
+        # on CONFLICTING_TOOL_PREFIX for more details
+        return CONFLICTING_TOOL_PREFIX + name
+
     return name
 
 
 def _denormalize_tool_name(name: str) -> str:
-    return name.removeprefix(CONFLICTING_TOOL_PREFIX)
+    if name.startswith(RESERVED_LC_TOOL_PREFIX):
+        assert "-" in name, "Invalid prefix in tool name"
+        _prefix, name = name.split("-", maxsplit=1)
+
+    return name
 
 
 def _agent_as_tool(agent: BaseAgent[OutputT]) -> StructuredTool:
@@ -757,17 +788,22 @@ async def _run(**kwargs: dict[str, Any]) -> OutputT | str:
 
 
 def _map_tool_call_from_langchain(tool_call: LC_ToolCall) -> ToolCall | SubagentCall:
-    if tool_call["name"].startswith(AGENT_PREFIX):
+    name = tool_call["name"]
+    if name.startswith(AGENT_PREFIX):
         return SubagentCall(
-            name=_denormalize_agent_name(tool_call["name"]),
+            name=_denormalize_agent_name(name),
             args=tool_call["args"],
             id=tool_call["id"],
         )
 
+    tool_type: ToolType = (
+        ToolType.LOCAL if name.startswith(LOCAL_TOOL_PREFIX) else ToolType.REMOTE
+    )
     return ToolCall(
-        name=_denormalize_tool_name(tool_call["name"]),
+        name=_denormalize_tool_name(name),
         args=tool_call["args"],
         id=tool_call["id"],
+        type=tool_type,
     )
 
 
@@ -776,13 +812,9 @@ def _map_tool_call_to_langchain(call: ToolCall | SubagentCall) -> LC_ToolCall:
         case SubagentCall():
             name = _normalize_agent_name(call.name)
         case ToolCall():
-            name = _normalize_tool_name(call.name)
+            name = _normalize_tool_name(call.name, call.type)
 
-    return LC_ToolCall(
-        name=name,
-        args=call.args,
-        id=call.id,
-    )
+    return LC_ToolCall(id=call.id, name=name, args=call.args)
 
 
 def _map_message_from_langchain(message: LC_BaseMessage) -> BaseMessage:
@@ -806,7 +838,7 @@ def _map_message_to_langchain(message: BaseMessage) -> LC_AnyMessage:
     match message:
         case AIMessage():
             lc_message = LC_AIMessage(content=message.content)
-            # this field can't be set via constructor
+            # This field can't be set via constructor
             lc_message.tool_calls = [
                 _map_tool_call_to_langchain(c) for c in message.calls
             ]

splunklib/ai/engines/langchain.py

@@ -12,21 +12,21 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import Any, Generic, Literal, TypeVar
 
 from pydantic import BaseModel
 
-OutputT = TypeVar("OutputT", default=None, covariant=True, bound=BaseModel | None)
+from splunklib.ai.tools import ToolType
 
 
 @dataclass(frozen=True)
 class ToolCall:
     name: str
     args: dict[str, Any]
     id: str | None  # TODO: can be None?
+    type: ToolType
 
 
 @dataclass(frozen=True)
@@ -41,7 +41,7 @@ class BaseMessage:
     role: str = ""
     content: str = field(default="")
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if type(self) is BaseMessage:
             raise TypeError(
                 "BaseMessage is an abstract class and cannot be instantiated"
@@ -79,14 +79,15 @@ class AIMessage(BaseMessage):
 
 @dataclass(frozen=True)
 class ToolMessage(BaseMessage):
-    """
-    ToolMessage represents a response of a tool call
-    """
+    """ToolMessage represents a response of a tool call"""
+
+    # TODO: See if we can remove the defaults - they should always be populated manually
 
     role: Literal["tool"] = "tool"
     name: str = field(default="")
     call_id: str = field(default="")
     status: Literal["success", "error"] = "success"
+    type: ToolType = ToolType.LOCAL
 
 
 @dataclass(frozen=True)
@@ -110,6 +111,9 @@ class SubagentMessage(BaseMessage):
     status: Literal["success", "error"] = "success"
 
 
+OutputT = TypeVar("OutputT", default=None, covariant=True, bound=BaseModel | None)
+
+
 @dataclass(frozen=True)
 class AgentResponse(Generic[OutputT]):
     # in case output_schema is provided, this will hold the parsed structured output