|
| 1 | +"""ADK-based agent using MCP protocol for tool access. |
| 2 | +
|
| 3 | +This module provides an ADK agent that can be used both in the ADK web UI |
| 4 | +and directly from the command line. The agent uses MCP tools to access |
| 5 | +external functionality. |
| 6 | +""" |
| 7 | + |
1 | 8 | import asyncio |
2 | 9 | import os |
3 | 10 | import logfire |
| 11 | +from typing import List, Tuple, Any |
4 | 12 |
|
5 | 13 | from dotenv import load_dotenv |
6 | 14 | from google.adk.agents.llm_agent import LlmAgent |
7 | | -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters |
| 15 | +from google.adk.tools.mcp_tool.mcp_toolset import ( |
| 16 | + MCPToolset, |
| 17 | + StdioServerParameters, |
| 18 | + StdioConnectionParams, |
| 19 | +) |
8 | 20 | from google.adk.runners import Runner |
9 | 21 | from google.adk.sessions import InMemorySessionService |
10 | 22 | from google.genai import types |
11 | 23 |
|
| 24 | +from agents_mcp_usage.utils import get_mcp_server_path |
| 25 | + |
12 | 26 | load_dotenv() |
13 | 27 |
|
14 | 28 | # Set API key for Google AI API from environment variable |
|
18 | 32 | logfire.configure(send_to_logfire="if-token-present", service_name="adk-basic-mcp") |
19 | 33 | logfire.instrument_mcp() |
20 | 34 |
|
| 35 | +# Global variable to store toolset instances for cleanup |
| 36 | +_ACTIVE_TOOLSETS = [] |
21 | 37 |
|
22 | | -async def main(query: str = "Greet Andrew and give him the current time") -> None: |
23 | | - """Runs the agent with a given query. |
24 | 38 |
|
25 | | - This function sets up the MCP server, creates an LLM agent, and runs it |
26 | | - with a specified query. It also handles the cleanup of the MCP server |
27 | | - connection. |
| 39 | +# Initialize the MCP tools and agent for ADK web UI |
| 40 | +async def init_mcp_toolset() -> Tuple[List[Any], List[MCPToolset]]: |
| 41 | + """Initialize and return the MCP tools and toolsets. |
28 | 42 |
|
29 | | - Args: |
30 | | - query: The query to run the agent with. |
| 43 | + Returns: |
| 44 | + Tuple[List[Any], List[MCPToolset]]: A tuple of (tools, toolsets) |
31 | 45 | """ |
32 | | - # Set up MCP server connection |
| 46 | + global _ACTIVE_TOOLSETS |
| 47 | + |
| 48 | + # Use absolute path to MCP server based on project root |
| 49 | + mcp_server_path = get_mcp_server_path("example_server.py") |
| 50 | + |
33 | 51 | server_params = StdioServerParameters( |
34 | 52 | command="uv", |
35 | | - args=["run", "mcp_servers/example_server.py", "stdio"], |
| 53 | + args=["run", str(mcp_server_path), "stdio"], |
36 | 54 | ) |
| 55 | + connection_params = StdioConnectionParams(server_params=server_params) |
| 56 | + toolset = MCPToolset(connection_params=connection_params) |
37 | 57 |
|
38 | | - tools, exit_stack = await MCPToolset.from_server(connection_params=server_params) |
39 | | - print(f"Connected to MCP server. Found {len(tools)} tools.") |
| 58 | + try: |
| 59 | + tools = await toolset.get_tools() |
| 60 | + _ACTIVE_TOOLSETS.append(toolset) |
| 61 | + return tools, [toolset] |
| 62 | + except Exception as e: |
| 63 | + # Clean up in case of initialization error |
| 64 | + try: |
| 65 | + await toolset.close() |
| 66 | + except Exception: |
| 67 | + pass |
| 68 | + raise e |
| 69 | + |
| 70 | + |
| 71 | +async def cleanup_toolsets(): |
| 72 | + """Clean up any active MCP toolset connections.""" |
| 73 | + global _ACTIVE_TOOLSETS |
| 74 | + |
| 75 | + for toolset in _ACTIVE_TOOLSETS: |
| 76 | + try: |
| 77 | + await toolset.close() |
| 78 | + print("MCP toolset connection closed.") |
| 79 | + except asyncio.CancelledError: |
| 80 | + print("MCP cleanup cancelled - this is normal") |
| 81 | + except Exception as e: |
| 82 | + print(f"Warning: Error during toolset cleanup: {e}") |
| 83 | + |
| 84 | + _ACTIVE_TOOLSETS = [] |
| 85 | + |
| 86 | + |
| 87 | +# Define a before_agent_callback to attach tools |
| 88 | +async def attach_tools_callback(callback_context): |
| 89 | + """Callback to attach tools to the agent before it runs. |
40 | 90 |
|
41 | | - # Create the agent |
42 | | - root_agent = LlmAgent( |
43 | | - model="gemini-2.5-pro-preview-03-25", |
44 | | - name="mcp_adk_assistant", |
45 | | - tools=tools, |
46 | | - ) |
| 91 | + Args: |
| 92 | + callback_context: The callback context from ADK. |
47 | 93 |
|
48 | | - # Set up session |
49 | | - session_service = InMemorySessionService() |
50 | | - session = session_service.create_session( |
51 | | - app_name="mcp_adk_app", |
52 | | - user_id="aginns", |
53 | | - ) |
| 94 | + Returns: |
| 95 | + None: The callback doesn't modify the content. |
| 96 | + """ |
| 97 | + await ensure_tools_attached() |
| 98 | + return None |
| 99 | + |
| 100 | + |
| 101 | +# This is the agent that will be imported by the ADK web UI |
| 102 | +root_agent = LlmAgent( |
| 103 | + model="gemini-2.0-flash", |
| 104 | + name="mcp_adk_assistant", |
| 105 | + instruction="You are an assistant that uses MCP tools to help users.", |
| 106 | + before_agent_callback=attach_tools_callback, # This ensures tools are attached |
| 107 | +) |
54 | 108 |
|
55 | | - # Create the runner |
56 | | - runner = Runner( |
57 | | - app_name="mcp_adk_app", |
58 | | - agent=root_agent, |
59 | | - session_service=session_service, |
60 | | - ) |
61 | 109 |
|
62 | | - # Run the agent with a query |
63 | | - content = types.Content(role="user", parts=[types.Part(text=query)]) |
| 110 | +# Flag to track if tools have been attached |
| 111 | +TOOLS_ATTACHED = False |
64 | 112 |
|
65 | | - print("Running agent...") |
| 113 | + |
| 114 | +# Function to dynamically attach tools to the agent |
| 115 | +async def ensure_tools_attached(): |
| 116 | + """Ensures that tools are attached to the agent before it's used.""" |
| 117 | + global TOOLS_ATTACHED |
| 118 | + |
| 119 | + if not TOOLS_ATTACHED: |
| 120 | + try: |
| 121 | + tools, _ = await init_mcp_toolset() |
| 122 | + print(f"✓ Connected to MCP server. Found {len(tools)} tools.") |
| 123 | + # Update the agent's tools |
| 124 | + root_agent.tools = tools |
| 125 | + TOOLS_ATTACHED = True |
| 126 | + except Exception as e: |
| 127 | + print(f"Error attaching MCP tools: {e}") |
| 128 | + # Set empty tools to avoid errors |
| 129 | + root_agent.tools = [] |
| 130 | + |
| 131 | + |
| 132 | +async def main(query: str = "Greet Andrew and give him the current time") -> None: |
| 133 | + """Runs the agent with a given query. |
| 134 | +
|
| 135 | + This function sets up a runner for the agent and runs it with a specified query. |
| 136 | + It also handles the cleanup of the MCP server connection. |
| 137 | +
|
| 138 | + Args: |
| 139 | + query: The query to run the agent with. |
| 140 | + """ |
66 | 141 | try: |
| 142 | + # Ensure tools are attached to the agent |
| 143 | + await ensure_tools_attached() |
| 144 | + |
| 145 | + # Set up session with async service |
| 146 | + session_service = InMemorySessionService() |
| 147 | + session = await session_service.create_session( |
| 148 | + app_name="mcp_adk_app", |
| 149 | + user_id="aginns", |
| 150 | + ) |
| 151 | + |
| 152 | + # Create the runner using the globally defined agent |
| 153 | + runner = Runner( |
| 154 | + app_name="mcp_adk_app", |
| 155 | + agent=root_agent, |
| 156 | + session_service=session_service, |
| 157 | + ) |
| 158 | + |
| 159 | + # Format the query as content |
| 160 | + content = types.Content(role="user", parts=[types.Part(text=query)]) |
| 161 | + |
| 162 | + print("Running agent...") |
67 | 163 | events_async = runner.run_async( |
68 | 164 | session_id=session.id, user_id=session.user_id, new_message=content |
69 | 165 | ) |
70 | 166 |
|
71 | 167 | async for event in events_async: |
72 | 168 | print(f"Event received: {event}") |
| 169 | + |
| 170 | + except Exception as e: |
| 171 | + print(f"Error during agent execution: {e}") |
| 172 | + print(f"Error type: {type(e).__name__}") |
| 173 | + raise |
73 | 174 | finally: |
74 | | - print("Closing MCP server connection...") |
75 | | - await exit_stack.aclose() |
76 | | - print("Cleanup complete.") |
| 175 | + # Clean up MCP toolsets to prevent asyncio shutdown errors |
| 176 | + await cleanup_toolsets() |
| 177 | + print("Agent execution completed successfully.") |
77 | 178 |
|
78 | 179 |
|
79 | 180 | if __name__ == "__main__": |
|
0 commit comments