108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
#!/usr/bin/env python
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
|
|
from app.agent.mcp import MCPAgent
|
|
from app.config import config
|
|
from app.logger import logger
|
|
|
|
|
|
class MCPRunner:
|
|
"""Runner class for MCP Agent with proper path handling and configuration."""
|
|
|
|
def __init__(self):
|
|
self.root_path = config.root_path
|
|
self.server_script = self.root_path / "app" / "mcp" / "server.py"
|
|
self.agent = MCPAgent()
|
|
|
|
async def initialize(self, connection_type: str, server_url: str = None) -> None:
|
|
"""Initialize the MCP agent with the appropriate connection."""
|
|
logger.info(f"Initializing MCPAgent with {connection_type} connection...")
|
|
|
|
if connection_type == "stdio":
|
|
await self.agent.initialize(
|
|
connection_type="stdio",
|
|
command=sys.executable,
|
|
args=[str(self.server_script)],
|
|
)
|
|
else: # sse
|
|
await self.agent.initialize(connection_type="sse", server_url=server_url)
|
|
|
|
logger.info(f"Connected to MCP server via {connection_type}")
|
|
|
|
async def run_interactive(self) -> None:
|
|
"""Run the agent in interactive mode."""
|
|
print("\nMCP Agent Interactive Mode (type 'exit' to quit)\n")
|
|
while True:
|
|
user_input = input("\nEnter your request: ")
|
|
if user_input.lower() in ["exit", "quit", "q"]:
|
|
break
|
|
response = await self.agent.run(user_input)
|
|
print(f"\nAgent: {response}")
|
|
|
|
async def run_single_prompt(self, prompt: str) -> None:
|
|
"""Run the agent with a single prompt."""
|
|
await self.agent.run(prompt)
|
|
|
|
async def run_default(self) -> None:
|
|
"""Run the agent in default mode."""
|
|
await self.agent.run(
|
|
"Hello, what tools are available to me? Terminate after you have listed the tools."
|
|
)
|
|
|
|
async def cleanup(self) -> None:
|
|
"""Clean up agent resources."""
|
|
await self.agent.cleanup()
|
|
logger.info("Session ended")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
"""Parse command line arguments."""
|
|
parser = argparse.ArgumentParser(description="Run the MCP Agent")
|
|
parser.add_argument(
|
|
"--connection",
|
|
"-c",
|
|
choices=["stdio", "sse"],
|
|
default="stdio",
|
|
help="Connection type: stdio or sse",
|
|
)
|
|
parser.add_argument(
|
|
"--server-url",
|
|
default="http://127.0.0.1:8000/sse",
|
|
help="URL for SSE connection",
|
|
)
|
|
parser.add_argument(
|
|
"--interactive", "-i", action="store_true", help="Run in interactive mode"
|
|
)
|
|
parser.add_argument("--prompt", "-p", help="Single prompt to execute and exit")
|
|
return parser.parse_args()
|
|
|
|
|
|
async def run_mcp() -> None:
|
|
"""Main entry point for the MCP runner."""
|
|
args = parse_args()
|
|
runner = MCPRunner()
|
|
|
|
try:
|
|
await runner.initialize(args.connection, args.server_url)
|
|
|
|
if args.prompt:
|
|
await runner.run_single_prompt(args.prompt)
|
|
elif args.interactive:
|
|
await runner.run_interactive()
|
|
else:
|
|
await runner.run_default()
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("Program interrupted by user")
|
|
except Exception as e:
|
|
logger.error(f"Error running MCPAgent: {str(e)}", exc_info=True)
|
|
sys.exit(1)
|
|
finally:
|
|
await runner.cleanup()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(run_mcp())
|