From 13a0316729a7298a60c2e4a092a7ad06f5196627 Mon Sep 17 00:00:00 2001 From: goulustis Date: Mon, 26 Jan 2026 16:22:30 +0800 Subject: [PATCH] choose tool set --- lang_agent/components/client_tool_manager.py | 24 ++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/lang_agent/components/client_tool_manager.py b/lang_agent/components/client_tool_manager.py index be9eb06..3f5cb2d 100644 --- a/lang_agent/components/client_tool_manager.py +++ b/lang_agent/components/client_tool_manager.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field, is_dataclass -from typing import Type, Any, Optional +from typing import Type, Any, Optional, List import tyro import commentjson import asyncio @@ -235,6 +235,9 @@ class ClientToolManagerConfig(InstantiateConfig): mcp_config_f: str = osp.join(osp.dirname(osp.dirname(osp.dirname(__file__))), "configs", "mcp_config.json") """path to all mcp configurations; expect json file""" + tool_keys: List = None + """tool configs to use; the keys inside mcp_config; if None, use everything""" + def __post_init__(self): assert osp.exists(self.mcp_config_f), f"config_f {self.mcp_config_f} does not exist." @@ -247,16 +250,29 @@ class ClientToolManager: def populate_module(self): with open(self.config.mcp_config_f, "r") as f: - self.mcp_configs = commentjson.load(f) + self.mcp_configs:dict = commentjson.load(f) async def aget_tools(self): """ Get tools from all configured MCP servers. Handles connection failures gracefully by logging warnings and continuing. """ - all_tools = [] - for server_name, server_config in self.mcp_configs.items(): + def get_to_load_configs() -> dict: + if self.config.tool_keys is None: + to_load_config = self.mcp_configs + else: + to_load_config = {} + for key in self.config.tool_keys: + val = self.mcp_configs.get(key) + if val is None: + logger.warning(f"{key} is not in mcp tools") + else: + to_load_config[key] = val + + to_load_config = get_to_load_configs() + all_tools = [] + for server_name, server_config in to_load_config.items(): try: # Create a client for this single server single_server_config = {server_name: server_config}