From 88032cf5c56bb2d81842353605f9f5ab4b2206ff Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Tue, 4 Nov 2025 15:47:44 -0800 Subject: [PATCH] feat: Support MCP prompts Add support for MCP prompts via the McpInstructionProvider class, which can be specified as an agent's instruction. Co-authored-by: Kathy Wu PiperOrigin-RevId: 828166051 --- contributing/samples/mcp_sse_agent/agent.py | 20 +- .../mcp_sse_agent/filesystem_server.py | 7 + src/google/adk/agents/__init__.py | 16 ++ .../adk/agents/mcp_instruction_provider.py | 107 +++++++++ .../agents/test_mcp_instruction_provider.py | 210 ++++++++++++++++++ 5 files changed, 351 insertions(+), 9 deletions(-) create mode 100644 src/google/adk/agents/mcp_instruction_provider.py create mode 100644 tests/unittests/agents/test_mcp_instruction_provider.py diff --git a/contributing/samples/mcp_sse_agent/agent.py b/contributing/samples/mcp_sse_agent/agent.py index 4981512e..8d0980df 100755 --- a/contributing/samples/mcp_sse_agent/agent.py +++ b/contributing/samples/mcp_sse_agent/agent.py @@ -16,25 +16,27 @@ import os from google.adk.agents.llm_agent import LlmAgent +from google.adk.agents.mcp_instruction_provider import McpInstructionProvider from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset _allowed_path = os.path.dirname(os.path.abspath(__file__)) +connection_params = SseConnectionParams( + url='http://localhost:3000/sse', + headers={'Accept': 'text/event-stream'}, +) + root_agent = LlmAgent( model='gemini-2.0-flash', name='enterprise_assistant', - instruction=f"""\ -Help user accessing their file systems. - -Allowed directory: {_allowed_path} - """, + instruction=McpInstructionProvider( + connection_params=connection_params, + prompt_name='file_system_prompt', + ), tools=[ MCPToolset( - connection_params=SseConnectionParams( - url='http://localhost:3000/sse', - headers={'Accept': 'text/event-stream'}, - ), + connection_params=connection_params, # don't want agent to do write operation # you can also do below # tool_filter=lambda tool, ctx=None: tool.name diff --git a/contributing/samples/mcp_sse_agent/filesystem_server.py b/contributing/samples/mcp_sse_agent/filesystem_server.py index cda4f0a9..291091e5 100644 --- a/contributing/samples/mcp_sse_agent/filesystem_server.py +++ b/contributing/samples/mcp_sse_agent/filesystem_server.py @@ -45,6 +45,13 @@ def get_cwd() -> str: return str(Path.cwd()) +# Add a prompt for accessing file systems +@mcp.prompt(name="file_system_prompt") +def file_system_prompt() -> str: + return f"""\ +Help the user access their file systems.""" + + # Graceful shutdown handler async def shutdown(signal, loop): """Cleanup tasks tied to the service's shutdown.""" diff --git a/src/google/adk/agents/__init__.py b/src/google/adk/agents/__init__.py index 498a9f5c..5710a21b 100644 --- a/src/google/adk/agents/__init__.py +++ b/src/google/adk/agents/__init__.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +import sys + from .base_agent import BaseAgent from .invocation_context import InvocationContext from .live_request_queue import LiveRequest @@ -35,3 +38,16 @@ __all__ = [ 'LiveRequestQueue', 'RunConfig', ] + +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_instruction_provider import McpInstructionProvider + + __all__.extend([ + 'McpInstructionProvider', + ]) diff --git a/src/google/adk/agents/mcp_instruction_provider.py b/src/google/adk/agents/mcp_instruction_provider.py new file mode 100644 index 00000000..e9f40663 --- /dev/null +++ b/src/google/adk/agents/mcp_instruction_provider.py @@ -0,0 +1,107 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Provides instructions to an agent by fetching prompts from an MCP server.""" + +from __future__ import annotations + +import logging +import sys +from typing import Any +from typing import Dict +from typing import TextIO + +from .llm_agent import InstructionProvider +from .readonly_context import ReadonlyContext + +# Attempt to import MCP Session Manager from the MCP library, and hints user to +# upgrade their Python version to 3.10 if it fails. +try: + from mcp import types + + from ..tools.mcp_tool.mcp_session_manager import MCPSessionManager +except ImportError as e: + if sys.version_info < (3, 10): + raise ImportError( + "MCP Session Manager requires Python 3.10 or above. Please upgrade" + " your Python version." + ) from e + else: + raise e + + +class McpInstructionProvider(InstructionProvider): + """Fetches agent instructions from an MCP server.""" + + def __init__( + self, + connection_params: Any, + prompt_name: str, + errlog: TextIO = sys.stderr, + ): + """Initializes the McpInstructionProvider. + + Args: + connection_params: Parameters for connecting to the MCP server. + prompt_name: The name of the MCP Prompt to fetch. + errlog: TextIO stream for error logging. + """ + self._connection_params = connection_params + self._errlog = errlog or logging.getLogger(__name__) + self._mcp_session_manager = MCPSessionManager( + connection_params=self._connection_params, + errlog=self._errlog, + ) + self.prompt_name = prompt_name + + async def __call__(self, context: ReadonlyContext) -> str: + """Fetches the instruction from the MCP server. + + Args: + context: The read-only context of the agent. + + Returns: + The instruction string. + """ + session = await self._mcp_session_manager.create_session() + # Fetch prompt definition to get the required argument names + prompt_definitions = await session.list_prompts() + prompt_definition = next( + (p for p in prompt_definitions.prompts if p.name == self.prompt_name), + None, + ) + + # Fetch arguments from context state if the prompt requires them + prompt_args: Dict[str, Any] = {} + if prompt_definition and prompt_definition.arguments: + arg_names = {arg.name for arg in prompt_definition.arguments} + prompt_args = { + k: v for k, v in (context.state or {}).items() if k in arg_names + } + + # Fetch the specific prompt by name with arguments from context state + prompt_result: types.GetPromptResult = await session.get_prompt( + self.prompt_name, arguments=prompt_args + ) + + if prompt_result and prompt_result.messages: + # Concatenate content of all messages to form the instruction. + instruction = "".join( + message.content.text + for message in prompt_result.messages + if message.content.type == "text" + ) + return instruction + else: + raise ValueError(f"Failed to load MCP prompt '{self.prompt_name}'.") diff --git a/tests/unittests/agents/test_mcp_instruction_provider.py b/tests/unittests/agents/test_mcp_instruction_provider.py new file mode 100644 index 00000000..1f2d098c --- /dev/null +++ b/tests/unittests/agents/test_mcp_instruction_provider.py @@ -0,0 +1,210 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Unit tests for McpInstructionProvider.""" +import sys +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.agents.readonly_context import ReadonlyContext +import pytest + +# Skip all tests in this module if Python version is less than 3.10 +pytestmark = pytest.mark.skipif( + sys.version_info < (3, 10), + reason="MCP instruction provider requires Python 3.10+", +) + +# Import dependencies with version checking +try: + from google.adk.agents.mcp_instruction_provider import McpInstructionProvider +except ImportError as e: + if sys.version_info < (3, 10): + # Create dummy classes to prevent NameError during test collection + # Tests will be skipped anyway due to pytestmark + class DummyClass: + pass + + McpInstructionProvider = DummyClass + else: + raise e + + +class TestMcpInstructionProvider: + """Unit tests for McpInstructionProvider.""" + + def setup_method(self): + """Sets up the test environment.""" + self.connection_params = {"host": "localhost", "port": 8000} + self.prompt_name = "test_prompt" + self.mock_mcp_session_manager_cls = patch( + "google.adk.agents.mcp_instruction_provider.MCPSessionManager" + ).start() + self.mock_mcp_session_manager = ( + self.mock_mcp_session_manager_cls.return_value + ) + self.mock_session = MagicMock() + self.mock_session.list_prompts = AsyncMock() + self.mock_session.get_prompt = AsyncMock() + self.mock_mcp_session_manager.create_session = AsyncMock( + return_value=self.mock_session + ) + self.provider = McpInstructionProvider( + self.connection_params, self.prompt_name + ) + + @pytest.mark.asyncio + async def test_call_success_no_args(self): + """Tests __call__ with a prompt that has no arguments.""" + mock_prompt = MagicMock() + mock_prompt.name = self.prompt_name + mock_prompt.arguments = None + self.mock_session.list_prompts.return_value = MagicMock( + prompts=[mock_prompt] + ) + + mock_msg1 = MagicMock() + mock_msg1.content.type = "text" + mock_msg1.content.text = "instruction part 1. " + mock_msg2 = MagicMock() + mock_msg2.content.type = "text" + mock_msg2.content.text = "instruction part 2" + self.mock_session.get_prompt.return_value = MagicMock( + messages=[mock_msg1, mock_msg2] + ) + + mock_invocation_context = MagicMock() + mock_invocation_context.session.state = {} + context = ReadonlyContext(mock_invocation_context) + + # Call + instruction = await self.provider(context) + + # Assert + assert instruction == "instruction part 1. instruction part 2" + self.mock_session.get_prompt.assert_called_once_with( + self.prompt_name, arguments={} + ) + + @pytest.mark.asyncio + async def test_call_success_with_args(self): + """Tests __call__ with a prompt that has arguments.""" + mock_arg1 = MagicMock() + mock_arg1.name = "arg1" + mock_prompt = MagicMock() + mock_prompt.name = self.prompt_name + mock_prompt.arguments = [mock_arg1] + self.mock_session.list_prompts.return_value = MagicMock( + prompts=[mock_prompt] + ) + + mock_msg = MagicMock() + mock_msg.content.type = "text" + mock_msg.content.text = "instruction with arg1" + self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg]) + + mock_invocation_context = MagicMock() + mock_invocation_context.session.state = {"arg1": "value1", "arg2": "value2"} + context = ReadonlyContext(mock_invocation_context) + + instruction = await self.provider(context) + + assert instruction == "instruction with arg1" + self.mock_session.get_prompt.assert_called_once_with( + self.prompt_name, arguments={"arg1": "value1"} + ) + + @pytest.mark.asyncio + async def test_call_prompt_not_found_in_list_prompts(self): + """Tests __call__ when list_prompts doesn't return the prompt.""" + self.mock_session.list_prompts.return_value = MagicMock(prompts=[]) + + mock_msg = MagicMock() + mock_msg.content.type = "text" + mock_msg.content.text = "instruction" + self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg]) + + mock_invocation_context = MagicMock() + mock_invocation_context.session.state = {"arg1": "value1"} + context = ReadonlyContext(mock_invocation_context) + + instruction = await self.provider(context) + + assert instruction == "instruction" + self.mock_session.get_prompt.assert_called_once_with( + self.prompt_name, arguments={} + ) + + @pytest.mark.asyncio + async def test_call_get_prompt_returns_no_messages(self): + """Tests __call__ when get_prompt returns no messages.""" + # Setup mocks + self.mock_session.list_prompts.return_value = MagicMock(prompts=[]) + self.mock_session.get_prompt.return_value = MagicMock(messages=[]) + + mock_invocation_context = MagicMock() + mock_invocation_context.session.state = {} + context = ReadonlyContext(mock_invocation_context) + + # Call and assert + with pytest.raises( + ValueError, match="Failed to load MCP prompt 'test_prompt'." + ): + await self.provider(context) + + # Assert + self.mock_session.get_prompt.assert_called_once_with( + self.prompt_name, arguments={} + ) + + @pytest.mark.asyncio + async def test_call_ignore_non_text_messages(self): + """Tests __call__ ignores non-text messages.""" + # Setup mocks + mock_prompt = MagicMock() + mock_prompt.name = self.prompt_name + mock_prompt.arguments = None + self.mock_session.list_prompts.return_value = MagicMock( + prompts=[mock_prompt] + ) + + mock_msg1 = MagicMock() + mock_msg1.content.type = "text" + mock_msg1.content.text = "instruction part 1. " + + mock_msg2 = MagicMock() + mock_msg2.content.type = "image" + mock_msg2.content.text = "ignored" + + mock_msg3 = MagicMock() + mock_msg3.content.type = "text" + mock_msg3.content.text = "instruction part 2" + + self.mock_session.get_prompt.return_value = MagicMock( + messages=[mock_msg1, mock_msg2, mock_msg3] + ) + + mock_invocation_context = MagicMock() + mock_invocation_context.session.state = {} + context = ReadonlyContext(mock_invocation_context) + + # Call + instruction = await self.provider(context) + + # Assert + assert instruction == "instruction part 1. instruction part 2" + self.mock_session.get_prompt.assert_called_once_with( + self.prompt_name, arguments={} + )