forked from va/knowledge-search-mcp
Add SSE support
This commit is contained in:
31
agent.py
31
agent.py
@@ -1,6 +1,7 @@
|
||||
# ruff: noqa: INP001
|
||||
"""ADK agent that connects to the knowledge-search MCP server."""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
@@ -8,7 +9,10 @@ from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions import InMemorySessionService
|
||||
from google.adk.tools.mcp_tool import McpToolset
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
|
||||
from google.adk.tools.mcp_tool.mcp_session_manager import (
|
||||
SseConnectionParams,
|
||||
StdioConnectionParams,
|
||||
)
|
||||
from google.genai import types
|
||||
from mcp import StdioServerParameters
|
||||
|
||||
@@ -22,15 +26,32 @@ if location := os.environ.get("LOCATION"):
|
||||
SERVER_SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "main.py")
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Knowledge Search Agent")
|
||||
parser.add_argument(
|
||||
"--remote",
|
||||
metavar="URL",
|
||||
help="Connect to an already-running MCP server at this SSE URL "
|
||||
"(e.g. http://localhost:8080/sse). Without this flag the agent "
|
||||
"spawns the server as a subprocess.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
async def async_main() -> None:
|
||||
toolset = McpToolset(
|
||||
connection_params=StdioConnectionParams(
|
||||
args = _parse_args()
|
||||
|
||||
if args.remote:
|
||||
connection_params = SseConnectionParams(url=args.remote)
|
||||
else:
|
||||
connection_params = StdioConnectionParams(
|
||||
server_params=StdioServerParameters(
|
||||
command="uv",
|
||||
args=["run", "python", SERVER_SCRIPT],
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
toolset = McpToolset(connection_params=connection_params)
|
||||
|
||||
agent = LlmAgent(
|
||||
model="gemini-2.0-flash",
|
||||
|
||||
Reference in New Issue
Block a user