diff --git a/owl/run_mcp.py b/owl/run_mcp.py index 84c201d..4ae556b 100644 --- a/owl/run_mcp.py +++ b/owl/run_mcp.py @@ -79,10 +79,10 @@ from camel.models import ModelFactory from camel.toolkits import FunctionTool from camel.types import ModelPlatformType, ModelType from camel.logger import set_log_level +from camel.toolkits import MCPToolkit from utils.enhanced_role_playing import OwlRolePlaying, run_society -from utils.mcp.mcp_toolkit_manager import MCPToolkitManager load_dotenv() @@ -138,7 +138,7 @@ async def main(): Path(__file__).parent / "utils/mcp/mcp_servers_config.json" ) - manager = MCPToolkitManager.from_config(config_path) + mcp_toolkit = MCPToolkit(config_path=config_path) question = ( "I'd like a academic report about Guohao Li, including his research " @@ -146,16 +146,18 @@ async def main(): "Then organize the report in Markdown format and save it to my desktop" ) - # Connect to all MCP toolkits - async with manager.connection(): - tools = manager.get_all_tools() + await mcp_toolkit.connect() - society = await construct_society(question, tools) + # # Connect to all MCP toolkits + tools = [*mcp_toolkit.get_tools()] - answer, chat_history, token_count = await run_society(society) + society = await construct_society(question, tools) + + answer, chat_history, token_count = await run_society(society) print(f"\033[94mAnswer: {answer}\033[0m") + await mcp_toolkit.disconnect() if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file diff --git a/owl/utils/mcp/mcp_toolkit_manager.py b/owl/utils/mcp/mcp_toolkit_manager.py deleted file mode 100644 index b7ceff6..0000000 --- a/owl/utils/mcp/mcp_toolkit_manager.py +++ /dev/null @@ -1,86 +0,0 @@ -import json -import os -from typing import List, Optional, AsyncGenerator - -from camel.toolkits import MCPToolkit -from contextlib import AsyncExitStack, asynccontextmanager - - -class MCPToolkitManager: - r"""MCPToolkitManager is a class for managing multiple MCPToolkit - instances and providing unified connection management. - - Attributes: - toolkits (List[MCPToolkit]): A list of MCPToolkit instances to be - managed. - """ - - def __init__(self, toolkits: List[MCPToolkit]): - self.toolkits = toolkits - self._exit_stack: Optional[AsyncExitStack] = None - self._connected = False - - - @staticmethod - def from_config(config_path: str) -> "MCPToolkitManager": - r"""Loads an MCPToolkit instance from a JSON configuration file and - returns an MCPToolkitManager instance. - - Args: - config_path (str): The path to the JSON configuration file. - - Returns: - MCPToolkitManager: The MCPToolkitManager instance. - """ - with open(config_path, "r", encoding="utf-8") as f: - data = json.load(f) - - all_toolkits = [] - - # "mcpServers" is the MCP server configuration running as stdio mode - mcp_servers = data.get("mcpServers", {}) - for name, cfg in mcp_servers.items(): - toolkit = MCPToolkit( - command_or_url=cfg["command"], - args=cfg.get("args", []), - env={**os.environ, **cfg.get("env", {})}, - timeout=cfg.get("timeout", None), - ) - all_toolkits.append(toolkit) - - # "mcpWebServers" is the MCP server configuration running as sse mode - mcp_web_servers = data.get("mcpWebServers", {}) - for name, cfg in mcp_web_servers.items(): - toolkit = MCPToolkit( - command_or_url=cfg["url"], - timeout=cfg.get("timeout", None), - ) - all_toolkits.append(toolkit) - - return MCPToolkitManager(all_toolkits) - - @asynccontextmanager - async def connection(self) -> AsyncGenerator["MCPToolkitManager", None]: - r"""Connect multiple MCPToolkit instances and close them when - leaving""" - self._exit_stack = AsyncExitStack() - try: - for tk in self.toolkits: - await self._exit_stack.enter_async_context(tk.connection()) - self._connected = True - yield self - finally: - self._connected = False - await self._exit_stack.aclose() - self._exit_stack = None - - def is_connected(self) -> bool: - r"""Returns whether the MCPToolkitManager is connected.""" - return self._connected - - def get_all_tools(self): - r"""Returns all tools from all MCPToolkit instances.""" - all_tools = [] - for tk in self.toolkits: - all_tools.extend(tk.get_tools()) - return all_tools \ No newline at end of file