mcp-use 1.1.5__py3-none-any.whl → 1.2.6__py3-none-any.whl

mcp_use/agents/mcpagent.py

@@ -5,16 +5,29 @@ This module provides the main MCPAgent class that integrates all components
 to provide a simple interface for using MCP tools with different LLMs.
 """
 
+import logging
+
+from langchain.agents import AgentExecutor, create_tool_calling_agent
+from langchain.globals import set_debug
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from langchain.schema.language_model import BaseLanguageModel
+from langchain_core.agents import AgentAction, AgentFinish
+from langchain_core.exceptions import OutputParserException
+from langchain_core.tools import BaseTool
+from langchain_core.utils.input import get_color_mapping
 
 from mcp_use.client import MCPClient
 from mcp_use.connectors.base import BaseConnector
 from mcp_use.session import MCPSession
 
+from ..adapters.langchain_adapter import LangChainAdapter
 from ..logging import logger
-from .langchain_agent import LangChainAgent
-from .prompts.default import DEFAULT_SYSTEM_PROMPT_TEMPLATE
+from .prompts.system_prompt_builder import create_system_message
+from .prompts.templates import DEFAULT_SYSTEM_PROMPT_TEMPLATE, SERVER_MANAGER_SYSTEM_PROMPT_TEMPLATE
+from .server_manager import ServerManager
+
+set_debug(logger.level == logging.DEBUG)
 
 
 class MCPAgent:
@@ -24,9 +37,6 @@ class MCPAgent:
     through LangChain's agent framework, with customizable system prompts and conversation memory.
     """
 
-    # Default system prompt template to use if none is provided
-    DEFAULT_SYSTEM_PROMPT_TEMPLATE = DEFAULT_SYSTEM_PROMPT_TEMPLATE
-
     def __init__(
         self,
         llm: BaseLanguageModel,
@@ -37,9 +47,11 @@ class MCPAgent:
         auto_initialize: bool = False,
         memory_enabled: bool = True,
         system_prompt: str | None = None,
-        system_prompt_template: str | None = None,
+        system_prompt_template: str | None = None,  # User can still override the template
         additional_instructions: str | None = None,
         disallowed_tools: list[str] | None = None,
+        use_server_manager: bool = False,
+        verbose: bool = False,
     ):
         """Initialize a new MCPAgent instance.
 
@@ -55,10 +67,11 @@ class MCPAgent:
             system_prompt_template: Template for system prompt with {tool_descriptions} placeholder.
             additional_instructions: Extra instructions to append to the system prompt.
             disallowed_tools: List of tool names that should not be available to the agent.
+            use_server_manager: Whether to use server manager mode instead of exposing all tools.
         """
         self.llm = llm
         self.client = client
-        self.connectors = connectors
+        self.connectors = connectors or []
         self.server_name = server_name
         self.max_steps = max_steps
         self.auto_initialize = auto_initialize
@@ -66,91 +79,148 @@ class MCPAgent:
         self._initialized = False
         self._conversation_history: list[BaseMessage] = []
         self.disallowed_tools = disallowed_tools or []
-
+        self.use_server_manager = use_server_manager
+        self.verbose = verbose
         # System prompt configuration
-        self.system_prompt = system_prompt
-        self.system_prompt_template = system_prompt_template or self.DEFAULT_SYSTEM_PROMPT_TEMPLATE
+        self.system_prompt = system_prompt  # User-provided full prompt override
+        # User can provide a template override, otherwise use the imported default
+        self.system_prompt_template_override = system_prompt_template
         self.additional_instructions = additional_instructions
 
         # Either client or connector must be provided
-        if not client and len(connectors) == 0:
+        if not client and len(self.connectors) == 0:
             raise ValueError("Either client or connector must be provided")
 
-        self._agent: LangChainAgent | None = None
-        self._sessions: dict[str, MCPSession] | None = None
+        # Create the adapter for tool conversion
+        self.adapter = LangChainAdapter(disallowed_tools=self.disallowed_tools)
+
+        # Initialize server manager if requested
+        self.server_manager = None
+        if self.use_server_manager:
+            if not self.client:
+                raise ValueError("Client must be provided when using server manager")
+            self.server_manager = ServerManager(self.client, self.adapter)
+
+        # State tracking
+        self._agent_executor: AgentExecutor | None = None
+        self._sessions: dict[str, MCPSession] = {}
         self._system_message: SystemMessage | None = None
+        self._tools: list[BaseTool] = []
 
     async def initialize(self) -> None:
         """Initialize the MCP client and agent."""
-        # If using client, get or create a session
-        if self.client:
-            # First try to get existing sessions
-            self._sessions = self.client.get_all_active_sessions()
-
-            # If no active sessions exist, create new ones
-            if not self._sessions:
-                self._sessions = await self.client.create_all_sessions()
-            connectors_to_use = [session.connector for session in self._sessions.values()]
+        logger.info("🚀 Initializing MCP agent and connecting to services...")
+        # If using server manager, initialize it
+        if self.use_server_manager and self.server_manager:
+            await self.server_manager.initialize()
+            # Get server management tools
+            management_tools = await self.server_manager.get_server_management_tools()
+            self._tools = management_tools
+            logger.info(
+                f"🔧 Server manager mode active with {len(management_tools)} management tools"
+            )
+
+            # Create the system message based on available tools
+            await self._create_system_message_from_tools(self._tools)
         else:
-            # Using direct connector
-            connectors_to_use = self.connectors
-            await [c_to_use.connect() for c_to_use in connectors_to_use]
-            await [c_to_use.initialize() for c_to_use in connectors_to_use]
-        # Create the system message based on available tools
-        await self._create_system_message(connectors_to_use)
+            # Standard initialization - if using client, get or create sessions
+            if self.client:
+                # First try to get existing sessions
+                self._sessions = self.client.get_all_active_sessions()
+                logger.info(f"🔌 Found {len(self._sessions)} existing sessions")
+
+                # If no active sessions exist, create new ones
+                if not self._sessions:
+                    logger.info("🔄 No active sessions found, creating new ones...")
+                    self._sessions = await self.client.create_all_sessions()
+                    logger.info(f"✅ Created {len(self._sessions)} new sessions")
+
+                # Create LangChain tools directly from the client using the adapter
+                self._tools = await self.adapter.create_tools(self.client)
+                logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from client")
+            else:
+                # Using direct connector - only establish connection
+                # LangChainAdapter will handle initialization
+                connectors_to_use = self.connectors
+                logger.info(f"🔗 Connecting to {len(connectors_to_use)} direct connectors...")
+                for connector in connectors_to_use:
+                    if not hasattr(connector, "client") or connector.client is None:
+                        await connector.connect()
+
+                # Create LangChain tools using the adapter with connectors
+                self._tools = await self.adapter._create_tools_from_connectors(connectors_to_use)
+                logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from connectors")
+
+            # Get all tools for system message generation
+            all_tools = self._tools
+            logger.info(f"🧰 Found {len(all_tools)} tools across all connectors")
+
+            # Create the system message based on available tools
+            await self._create_system_message_from_tools(all_tools)
 
         # Create the agent
-        self._agent = LangChainAgent(
-            connectors=connectors_to_use,
-            llm=self.llm,
-            max_steps=self.max_steps,
-            system_message=(self._system_message.content if self._system_message else None),
+        self._agent_executor = self._create_agent()
+        self._initialized = True
+        logger.info("✨ Agent initialization complete")
+
+    async def _create_system_message_from_tools(self, tools: list[BaseTool]) -> None:
+        """Create the system message based on provided tools using the builder."""
+        # Use the override if provided, otherwise use the imported default
+        default_template = self.system_prompt_template_override or DEFAULT_SYSTEM_PROMPT_TEMPLATE
+        # Server manager template is now also imported
+        server_template = SERVER_MANAGER_SYSTEM_PROMPT_TEMPLATE
+
+        # Delegate creation to the imported function
+        self._system_message = create_system_message(
+            tools=tools,
+            system_prompt_template=default_template,
+            server_manager_template=server_template,  # Pass the imported template
+            use_server_manager=self.use_server_manager,
             disallowed_tools=self.disallowed_tools,
+            user_provided_prompt=self.system_prompt,
+            additional_instructions=self.additional_instructions,
         )
 
-        # Initialize the agent
-        await self._agent.initialize()
-        self._initialized = True
+        # Update conversation history if memory is enabled
+        if self.memory_enabled:
+            history_without_system = [
+                msg for msg in self._conversation_history if not isinstance(msg, SystemMessage)
+            ]
+            self._conversation_history = [self._system_message] + history_without_system
 
-    async def _create_system_message(self, connectors: list[BaseConnector]) -> None:
-        """Create the system message based on available tools.
+    def _create_agent(self) -> AgentExecutor:
+        """Create the LangChain agent with the configured system message.
 
-        Args:
-            connector: The connector with available tools.
+        Returns:
+            An initialized AgentExecutor.
         """
-        # If a complete system prompt was provided, use it directly
-        if self.system_prompt:
-            self._system_message = SystemMessage(content=self.system_prompt)
-            return
-
-        # Otherwise, build the system prompt from the template and tool descriptions
-        tool_descriptions = []
-        for connector in connectors:
-            tools = connector.tools
-            # Generate tool descriptions
-            for tool in tools:
-                # Skip disallowed tools
-                if tool.name in self.disallowed_tools:
-                    continue
-
-                # Escape curly braces in the description by doubling them
-                # (sometimes e.g. blender mcp they are used in the description)
-                description = (
-                    f"- {tool.name}: {tool.description.replace('{', '{{').replace('}', '}}')}"
-                )
-                tool_descriptions.append(description)
-
-        # Format the system prompt template with tool descriptions
-        system_prompt = self.system_prompt_template.format(
-            tool_descriptions="\n".join(tool_descriptions)
+        logger.debug(f"Creating new agent with {len(self._tools)} tools")
+
+        system_content = "You are a helpful assistant"
+        if self._system_message:
+            system_content = self._system_message.content
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", system_content),
+                MessagesPlaceholder(variable_name="chat_history"),
+                ("human", "{input}"),
+                MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ]
         )
 
-        # Add any additional instructions
-        if self.additional_instructions:
-            system_prompt += f"\n\n{self.additional_instructions}"
+        tool_names = [tool.name for tool in self._tools]
+        logger.info(f"🧠 Agent ready with tools: {', '.join(tool_names)}")
+
+        # Use the standard create_tool_calling_agent
+        agent = create_tool_calling_agent(llm=self.llm, tools=self._tools, prompt=prompt)
 
-        # Create the system message
-        self._system_message = SystemMessage(content=system_prompt)
+        # Use the standard AgentExecutor
+        executor = AgentExecutor(
+            agent=agent, tools=self._tools, max_iterations=self.max_steps, verbose=self.verbose
+        )
+        logger.debug(f"Created agent executor with max_iterations={self.max_steps}")
+        return executor
 
     def get_conversation_history(self) -> list[BaseMessage]:
         """Get the current conversation history.
@@ -165,7 +235,7 @@ class MCPAgent:
         self._conversation_history = []
 
         # Re-add the system message if it exists
-        if self._system_message:
+        if self._system_message and self.memory_enabled:
             self._conversation_history = [self._system_message]
 
     def add_to_history(self, message: BaseMessage) -> None:
@@ -174,7 +244,8 @@ class MCPAgent:
         Args:
             message: The message to add.
         """
-        self._conversation_history.append(message)
+        if self.memory_enabled:
+            self._conversation_history.append(message)
 
     def get_system_message(self) -> SystemMessage | None:
         """Get the current system message.
@@ -192,9 +263,21 @@ class MCPAgent:
         """
         self._system_message = SystemMessage(content=message)
 
-        # Update the agent if initialized
-        if self._agent:
-            self._agent.set_system_message(message)
+        # Update conversation history if memory is enabled
+        if self.memory_enabled:
+            # Remove old system message if it exists
+            history_without_system = [
+                msg for msg in self._conversation_history if not isinstance(msg, SystemMessage)
+            ]
+            self._conversation_history = history_without_system
+
+            # Add new system message
+            self._conversation_history.insert(0, self._system_message)
+
+        # Recreate the agent with the new system message if initialized
+        if self._initialized and self._tools:
+            self._agent_executor = self._create_agent()
+            logger.debug("Agent recreated with new system message")
 
     def set_disallowed_tools(self, disallowed_tools: list[str]) -> None:
         """Set the list of tools that should not be available to the agent.
@@ -205,6 +288,7 @@ class MCPAgent:
             disallowed_tools: List of tool names that should not be available.
         """
         self.disallowed_tools = disallowed_tools
+        self.adapter.disallowed_tools = disallowed_tools
 
         # If the agent is already initialized, we need to reinitialize it
         # to apply the changes to the available tools
@@ -230,7 +314,7 @@ class MCPAgent:
         manage_connector: bool = True,
         external_history: list[BaseMessage] | None = None,
     ) -> str:
-        """Run a query using the MCP tools.
+        """Run a query using the MCP tools with unified step-by-step execution.
 
         This method handles connecting to the MCP server, initializing the agent,
         running the query, and then cleaning up the connection.
@@ -253,19 +337,28 @@ class MCPAgent:
 
         try:
             # Initialize if needed
-            if manage_connector and (not self._initialized or not self._agent):
-                logger.debug("Initializing agent before running query")
+            if manage_connector and not self._initialized:
                 await self.initialize()
                 initialized_here = True
             elif not self._initialized and self.auto_initialize:
-                logger.debug("Auto-initializing agent before running query")
                 await self.initialize()
                 initialized_here = True
 
             # Check if initialization succeeded
-            if not self._agent:
+            if not self._agent_executor:
                 raise RuntimeError("MCP agent failed to initialize")
 
+            steps = max_steps or self.max_steps
+            if self._agent_executor:
+                self._agent_executor.max_iterations = steps
+
+            display_query = (
+                query[:50].replace("\n", " ") + "..."
+                if len(query) > 50
+                else query.replace("\n", " ")
+            )
+            logger.info(f"💬 Received query: '{display_query}'")
+
             # Add the user query to conversation history if memory is enabled
             if self.memory_enabled:
                 self.add_to_history(HumanMessage(content=query))
@@ -275,74 +368,164 @@ class MCPAgent:
                 external_history if external_history is not None else self._conversation_history
             )
 
-            # Convert messages to format expected by LangChain
+            # Convert messages to format expected by LangChain agent input
+            # Exclude the main system message as it's part of the agent's prompt
             langchain_history = []
             for msg in history_to_use:
                 if isinstance(msg, HumanMessage):
-                    langchain_history.append({"type": "human", "content": msg.content})
+                    langchain_history.append(msg)
                 elif isinstance(msg, AIMessage):
-                    langchain_history.append({"type": "ai", "content": msg.content})
-                elif isinstance(msg, SystemMessage) and msg != self._system_message:
-                    # Include system messages other than the main one
-                    # which is already included in the agent's prompt
-                    langchain_history.append({"type": "system", "content": msg.content})
-                # Other message types can be handled here if needed
-
-            # Run the query with the specified max_steps or default
-            logger.debug(f"Running query with max_steps={max_steps or self.max_steps}")
-            result = await self._agent.run(
-                query=query,
-                max_steps=max_steps,
-                chat_history=langchain_history,
+                    langchain_history.append(msg)
+
+            intermediate_steps: list[tuple[AgentAction, str]] = []
+            inputs = {"input": query, "chat_history": langchain_history}
+
+            # Construct a mapping of tool name to tool for easy lookup
+            name_to_tool_map = {tool.name: tool for tool in self._tools}
+            color_mapping = get_color_mapping(
+                [tool.name for tool in self._tools], excluded_colors=["green", "red"]
             )
 
-            # Add the response to conversation history if memory is enabled
+            logger.info(f"🏁 Starting agent execution with max_steps={steps}")
+
+            for step_num in range(steps):
+                # --- Check for tool updates if using server manager ---
+                if self.use_server_manager and self.server_manager:
+                    current_tools = await self.server_manager.get_all_tools()
+                    current_tool_names = {tool.name for tool in current_tools}
+                    existing_tool_names = {tool.name for tool in self._tools}
+
+                    if current_tool_names != existing_tool_names:
+                        logger.info(
+                            f"🔄 Tools changed before step {step_num + 1}, updating agent. "
+                            f"New tools: {', '.join(current_tool_names)}"
+                        )
+                        self._tools = current_tools
+                        # Regenerate system message with ALL current tools
+                        await self._create_system_message_from_tools(self._tools)
+                        # Recreate the agent executor with the new tools and system message
+                        self._agent_executor = self._create_agent()
+                        self._agent_executor.max_iterations = steps
+                        # Update maps for this iteration
+                        name_to_tool_map = {tool.name: tool for tool in self._tools}
+                        color_mapping = get_color_mapping(
+                            [tool.name for tool in self._tools], excluded_colors=["green", "red"]
+                        )
+
+                logger.info(f"🔍 Step {step_num + 1}/{steps}")
+
+                # --- Plan and execute the next step ---
+                try:
+                    # Use the internal _atake_next_step which handles planning and execution
+                    # This requires providing the necessary context like maps and intermediate steps
+                    next_step_output = await self._agent_executor._atake_next_step(
+                        name_to_tool_map=name_to_tool_map,
+                        color_mapping=color_mapping,
+                        inputs=inputs,
+                        intermediate_steps=intermediate_steps,
+                        run_manager=None,
+                    )
+
+                    # Process the output
+                    if isinstance(next_step_output, AgentFinish):
+                        logger.info(f"✅ Agent finished at step {step_num + 1}")
+                        result = next_step_output.return_values.get("output", "No output generated")
+                        break
+
+                    # If it's actions/steps, add to intermediate steps
+                    intermediate_steps.extend(next_step_output)
+
+                    # Log tool calls
+                    for action, output in next_step_output:
+                        tool_name = action.tool
+                        tool_input_str = str(action.tool_input)
+                        # Truncate long inputs for readability
+                        if len(tool_input_str) > 100:
+                            tool_input_str = tool_input_str[:97] + "..."
+                        logger.info(f"🔧 Tool call: {tool_name} with input: {tool_input_str}")
+                        # Truncate long outputs for readability
+                        output_str = str(output)
+                        if len(output_str) > 100:
+                            output_str = output_str[:97] + "..."
+                        output_str = output_str.replace("\n", " ")
+                        logger.info(f"📄 Tool result: {output_str}")
+
+                    # Check for return_direct on the last action taken
+                    if len(next_step_output) > 0:
+                        last_step: tuple[AgentAction, str] = next_step_output[-1]
+                        tool_return = self._agent_executor._get_tool_return(last_step)
+                        if tool_return is not None:
+                            logger.info(f"🏆 Tool returned directly at step {step_num + 1}")
+                            result = tool_return.return_values.get("output", "No output generated")
+                            break
+
+                except OutputParserException as e:
+                    logger.error(f"❌ Output parsing error during step {step_num + 1}: {e}")
+                    result = f"Agent stopped due to a parsing error: {str(e)}"
+                    break
+                except Exception as e:
+                    logger.error(f"❌ Error during agent execution step {step_num + 1}: {e}")
+                    import traceback
+
+                    traceback.print_exc()
+                    result = f"Agent stopped due to an error: {str(e)}"
+                    break
+
+            # --- Loop finished ---
+            if not result:
+                logger.warning(f"⚠️ Agent stopped after reaching max iterations ({steps})")
+                result = f"Agent stopped after reaching the maximum number of steps ({steps})."
+
+            # Add the final response to conversation history if memory is enabled
             if self.memory_enabled:
                 self.add_to_history(AIMessage(content=result))
 
+            logger.info("🎉 Agent execution complete")
             return result
 
         except Exception as e:
-            logger.error(f"Error running query: {e}")
-            # If we initialized in this method and there was an error,
-            # make sure to clean up
+            logger.error(f"❌ Error running query: {e}")
             if initialized_here and manage_connector:
-                logger.debug("Cleaning up resources after initialization error")
+                logger.info("🧹 Cleaning up resources after initialization error in run")
                 await self.close()
             raise
 
         finally:
-            # Clean up resources if we're managing the connector and
-            # we're not using a client that manages sessions
+            # Clean up if necessary (e.g., if not using client-managed sessions)
             if manage_connector and not self.client and not initialized_here:
-                logger.debug("Closing agent after query completion")
+                logger.info("🧹 Closing agent after query completion")
                 await self.close()
 
     async def close(self) -> None:
         """Close the MCP connection with improved error handling."""
+        logger.info("🔌 Closing agent and cleaning up resources...")
         try:
-            if self._agent:
-                # Clean up the agent first
-                logger.debug("Cleaning up agent")
-                self._agent = None
+            # Clean up the agent first
+            self._agent_executor = None
+            self._tools = []
 
             # If using client with session, close the session through client
-            if self.client and self._sessions:
-                logger.debug("Closing session through client")
+            if self.client:
+                logger.info("🔄 Closing sessions through client")
                 await self.client.close_all_sessions()
-                self._session = None
+                self._sessions = {}
             # If using direct connector, disconnect
             elif self.connectors:
                 for connector in self.connectors:
-                    logger.debug("Disconnecting connector")
+                    logger.info("🔄 Disconnecting connector")
                     await connector.disconnect()
 
+            # Clear adapter tool cache
+            if hasattr(self.adapter, "_connector_tool_map"):
+                self.adapter._connector_tool_map = {}
+
             self._initialized = False
-            logger.debug("Agent closed successfully")
+            logger.info("👋 Agent closed successfully")
 
         except Exception as e:
-            logger.error(f"Error during agent closure: {e}")
+            logger.error(f"❌ Error during agent closure: {e}")
             # Still try to clean up references even if there was an error
-            self._agent = None
-            self._session = None
+            self._agent_executor = None
+            self._tools = []
+            self._sessions = {}
             self._initialized = False