load custom toolsets

Author: twishabansal

packages/toolbox-core/src/mcp_poc/mcp_transport.py

@@ -1,10 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Optional, Dict
 import httpx
 
 class ITransport(ABC):
     @abstractmethod
-    async def tools_list(self) -> Dict:
+    async def tools_list(self, toolset_name: Optional[str] = None) -> Dict:
         pass
 
     @abstractmethod
@@ -25,16 +25,26 @@ def _build_json_rpc_payload(self, method: str, params: Dict) -> Dict:
         self._request_id += 1
         return {"jsonrpc": "2.0", "method": method, "params": params, "id": self._request_id}
 
-    async def tools_list(self) -> Dict:
+    def _get_list_endpoint(self, toolset_name: Optional[str] = None) -> str:
+        """Constructs the correct API endpoint for listing tools."""
+        if toolset_name:
+            return f"{self._base_url}/mcp/{toolset_name}"
+        return f"{self._base_url}/mcp"
+
+    async def tools_list(self, toolset_name: Optional[str] = None) -> Dict:
+        """Lists tools from the default endpoint or a specific toolset."""
+        endpoint = self._get_list_endpoint(toolset_name)
         payload = self._build_json_rpc_payload("tools/list", {})
-        response = await self._client.post(f"{self._base_url}/mcp", json=payload)
+        response = await self._client.post(endpoint, json=payload)
         response.raise_for_status()
         return response.json()
 
     async def tool_invoke(self, tool_name: str, arguments: Dict) -> Dict:
+        """Invokes a tool using the global /mcp endpoint."""
+        endpoint = f"{self._base_url}/mcp"
         params = {"name": tool_name, "arguments": arguments}
         payload = self._build_json_rpc_payload("tools/call", params)
-        response = await self._client.post(f"{self._base_url}/mcp", json=payload)
+        response = await self._client.post(endpoint, json=payload)
         response.raise_for_status()
         return response.json()
 

packages/toolbox-core/src/mcp_poc/run.py

@@ -1,20 +1,29 @@
 import asyncio
 import json
 from mcp_transport import McpHttpTransport
+from typing import Optional
 
 class MCPClient:
     """A simple client to interact with an MCP server."""
     def __init__(self, base_url: str):
         self._transport = McpHttpTransport(base_url=base_url)
 
-    async def list_all_tools(self):
-        print("--> Attempting to list tools...")
-        response = await self._transport.tools_list()
+    async def list_tools(self, toolset_name: Optional[str] = None):
+        """Lists tools, either all or from a specific toolset."""
+        if toolset_name:
+            print(f"--> Attempting to list tools from toolset: '{toolset_name}'...")
+        else:
+            print("--> Attempting to list all tools...")
+        
+        response = await self._transport.tools_list(toolset_name=toolset_name)
         return response.get("result", {}).get("tools", [])
 
-    async def invoke_a_tool(self, tool_name: str, args: dict):
+    async def invoke_tool(self, tool_name: str, args: dict):
+        """Invokes a tool using the global endpoint."""
         print(f"\n--> Attempting to invoke tool: '{tool_name}'...")
         response = await self._transport.tool_invoke(tool_name, args)
+
+        
         return response.get("result", {})
 
     async def close(self):
@@ -27,14 +36,20 @@ async def main():
 
     try:
         # 1. List all available tools
-        tools = await client.list_all_tools()
-        print("\n✅ Tools listed successfully:")
-        print(json.dumps(tools, indent=2))
+        all_tools = await client.list_tools()
+        print("\n✅ All tools listed successfully:")
+        print(json.dumps(all_tools, indent=2))
+
+        # 2. List tools from a specific toolset
+        custom_toolset_name = "my-toolset-2"
+        custom_tools = await client.list_tools(toolset_name=custom_toolset_name)
+        print(f"\n✅ Tools from '{custom_toolset_name}' toolset listed successfully:")
+        print(json.dumps(custom_tools, indent=2))
 
-        # 2. Invoke a specific tool
+        # 3. Invoke a tool. This correctly uses the global endpoint.
         tool_to_invoke = "get-n-rows"
         arguments = {"num_rows": "2"}
-        invocation_result = await client.invoke_a_tool(tool_to_invoke, arguments)
+        invocation_result = await client.invoke_tool(tool_to_invoke, arguments)
 
         print("\n✅ Tool invoked successfully:")
         print(json.dumps(invocation_result, indent=2))