You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
dc193f7969
commit
d9ce2e691c
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user