feat(config): implement config and from_config for MCPToolset

The connection_params argument in the constructor is split into four arguments in the config class because some of them have identical fields. In order to identify which is which, a separate name is more convenient.

PiperOrigin-RevId: 791965995
This commit is contained in:
Liang Wu
2025-08-06 20:45:22 -07:00
committed by Copybara-Service
parent dc193f7969
commit d9ce2e691c
2 changed files with 85 additions and 1 deletions
+16 -1
View File
@@ -11,7 +11,8 @@
# 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.
import logging
import sys
from ..auth.auth_tool import AuthToolArguments
from .agent_tool import AgentTool
@@ -52,3 +53,17 @@ __all__ = [
'ToolContext',
'transfer_to_agent',
]
if sys.version_info < (3, 10):
logger = logging.getLogger('google_adk.' + __name__)
logger.warning(
'MCP requires Python 3.10 or above. Please upgrade your Python'
' version in order to use it.'
)
else:
from .mcp_tool.mcp_toolset import MCPToolset
__all__.extend([
'MCPToolset',
])
@@ -21,12 +21,17 @@ from typing import Optional
from typing import TextIO
from typing import Union
from pydantic import model_validator
from typing_extensions import override
from ...agents.readonly_context import ReadonlyContext
from ...auth.auth_credential import AuthCredential
from ...auth.auth_schemes import AuthScheme
from ..base_tool import BaseTool
from ..base_toolset import BaseToolset
from ..base_toolset import ToolPredicate
from ..tool_configs import BaseToolConfig
from ..tool_configs import ToolArgsConfig
from .mcp_session_manager import MCPSessionManager
from .mcp_session_manager import retry_on_closed_resource
from .mcp_session_manager import SseConnectionParams
@@ -178,3 +183,67 @@ class MCPToolset(BaseToolset):
except Exception as e:
# Log the error but don't re-raise to avoid blocking shutdown
print(f"Warning: Error during MCPToolset cleanup: {e}", file=self._errlog)
@override
@classmethod
def from_config(
cls: type[MCPToolset], config: ToolArgsConfig, config_abs_path: str
) -> MCPToolset:
"""Creates an MCPToolset from a configuration object."""
mcp_toolset_config = MCPToolsetConfig.model_validate(config.model_dump())
if mcp_toolset_config.stdio_server_params:
connection_params = mcp_toolset_config.stdio_server_params
elif mcp_toolset_config.stdio_connection_params:
connection_params = mcp_toolset_config.stdio_connection_params
elif mcp_toolset_config.sse_connection_params:
connection_params = mcp_toolset_config.sse_connection_params
elif mcp_toolset_config.streamable_http_connection_params:
connection_params = mcp_toolset_config.streamable_http_connection_params
else:
raise ValueError("No connection params found in MCPToolsetConfig.")
return cls(
connection_params=connection_params,
tool_filter=mcp_toolset_config.tool_filter,
auth_scheme=mcp_toolset_config.auth_scheme,
auth_credential=mcp_toolset_config.auth_credential,
)
class MCPToolsetConfig(BaseToolConfig):
"""The config for MCPToolset."""
stdio_server_params: Optional[StdioServerParameters] = None
stdio_connection_params: Optional[StdioConnectionParams] = None
sse_connection_params: Optional[SseConnectionParams] = None
streamable_http_connection_params: Optional[
StreamableHTTPConnectionParams
] = None
tool_filter: Optional[List[str]] = None
auth_scheme: Optional[AuthScheme] = None
auth_credential: Optional[AuthCredential] = None
@model_validator(mode="after")
def _check_only_one_params_field(self):
param_fields = [
self.stdio_server_params,
self.stdio_connection_params,
self.sse_connection_params,
self.streamable_http_connection_params,
]
populated_fields = [f for f in param_fields if f is not None]
if len(populated_fields) != 1:
raise ValueError(
"Exactly one of stdio_server_params, stdio_connection_params,"
" sse_connection_params, streamable_http_connection_params must be"
" set."
)
return self