mcp-use 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

mcp_use/__init__.py

@@ -1,5 +1,5 @@
 """
-mcp_use - A model-agnostic MCP (Multi-Channel Platform) library for LLMs.
+mcp_use - An MCP library for LLMs.
 
 This library provides a unified interface for connecting different LLMs
 to MCP tools through existing LangChain adapters.

mcp_use/agents/langchain_agent.py

@@ -1,8 +1,8 @@
 """
-LangChain agent implementation for MCP tools.
+LangChain agent implementation for MCP tools with customizable system message.
 
 This module provides a LangChain agent implementation that can use MCP tools
-through a unified interface.
+through a unified interface, with support for customizable system messages.
 """
 
 from typing import Any, NoReturn
@@ -73,8 +73,15 @@ class LangChainAgent:
     through a unified interface.
     """
 
+    # Default system message if none is provided
+    DEFAULT_SYSTEM_MESSAGE = "You are a helpful AI assistant that can use tools to help users."
+
     def __init__(
-        self, connector: BaseConnector, llm: BaseLanguageModel, max_steps: int = 5
+        self,
+        connector: BaseConnector,
+        llm: BaseLanguageModel,
+        max_steps: int = 5,
+        system_message: str | None = None,
     ) -> None:
         """Initialize a new LangChain agent.
 
@@ -82,13 +89,28 @@ class LangChainAgent:
             connector: The MCP connector to use.
             llm: The LangChain LLM to use.
             max_steps: The maximum number of steps to take.
+            system_message: Optional custom system message to use.
         """
         self.connector = connector
         self.llm = llm
         self.max_steps = max_steps
+        self.system_message = system_message or self.DEFAULT_SYSTEM_MESSAGE
         self.tools: list[BaseTool] = []
         self.agent: AgentExecutor | None = None
 
+    def set_system_message(self, message: str) -> None:
+        """Set a new system message and recreate the agent.
+
+        Args:
+            message: The new system message.
+        """
+        self.system_message = message
+
+        # Recreate the agent with the new system message if it exists
+        if self.agent and self.tools:
+            self.agent = self._create_agent()
+            logger.info("Agent recreated with new system message")
+
     async def initialize(self) -> None:
         """Initialize the agent and its tools."""
         self.tools = await self._create_langchain_tools()
@@ -184,7 +206,7 @@ class LangChainAgent:
         return langchain_tools
 
     def _create_agent(self) -> AgentExecutor:
-        """Create the LangChain agent.
+        """Create the LangChain agent with the configured system message.
 
         Returns:
             An initialized AgentExecutor.
@@ -193,7 +215,7 @@ class LangChainAgent:
             [
                 (
                     "system",
-                    "You are a helpful AI assistant that can use tools to help users.",
+                    self.system_message,
                 ),
                 MessagesPlaceholder(variable_name="chat_history"),
                 ("human", "{input}"),
@@ -234,7 +256,25 @@ class LangChainAgent:
         if chat_history is None:
             chat_history = []
 
+        # Add a hint to use tools for queries about current information
+        enhanced_query = query
+        if any(
+            keyword in query.lower()
+            for keyword in [
+                "weather",
+                "current",
+                "today",
+                "now",
+                "latest",
+                "news",
+                "price",
+                "stock",
+            ]
+        ):
+            # Just log this, don't modify the query
+            logger.info("Query involves current information that may benefit from tool use")
+
         # Invoke with all required variables
-        result = await self.agent.ainvoke({"input": query, "chat_history": chat_history})
+        result = await self.agent.ainvoke({"input": enhanced_query, "chat_history": chat_history})
 
         return result["output"]

mcp_use/agents/mcpagent.py

@@ -1,10 +1,11 @@
 """
-Model-Agnostic MCP: Main integration module.
+MCP: Main integration module with customizable system prompt.
 
 This module provides the main MCPAgent class that integrates all components
 to provide a simple interface for using MCP tools with different LLMs.
 """
 
+from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from langchain.schema.language_model import BaseLanguageModel
 
 from mcp_use.client import MCPClient
@@ -13,15 +14,19 @@ from mcp_use.session import MCPSession
 
 from ..logging import logger
 from .langchain_agent import LangChainAgent
+from .prompts.default import DEFAULT_SYSTEM_PROMPT_TEMPLATE
 
 
 class MCPAgent:
     """Main class for using MCP tools with various LLM providers.
 
     This class provides a unified interface for using MCP tools with different LLM providers
-    through LangChain's agent framework.
+    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,
@@ -30,6 +35,10 @@ class MCPAgent:
         server_name: str | None = None,
         max_steps: int = 5,
         auto_initialize: bool = False,
+        memory_enabled: bool = True,
+        system_prompt: str | None = None,
+        system_prompt_template: str | None = None,
+        additional_instructions: str | None = None,
     ):
         """Initialize a new MCPAgent instance.
 
@@ -40,6 +49,10 @@ class MCPAgent:
             server_name: The name of the server to use if client is provided.
             max_steps: The maximum number of steps to take.
             auto_initialize: Whether to automatically initialize the agent when run is called.
+            memory_enabled: Whether to maintain conversation history for context.
+            system_prompt: Complete system prompt to use (overrides template if provided).
+            system_prompt_template: Template for system prompt with {tool_descriptions} placeholder.
+            additional_instructions: Extra instructions to append to the system prompt.
         """
         self.llm = llm
         self.client = client
@@ -47,7 +60,14 @@ class MCPAgent:
         self.server_name = server_name
         self.max_steps = max_steps
         self.auto_initialize = auto_initialize
+        self.memory_enabled = memory_enabled
         self._initialized = False
+        self._conversation_history: list[BaseMessage] = []
+
+        # System prompt configuration
+        self.system_prompt = system_prompt
+        self.system_prompt_template = system_prompt_template or self.DEFAULT_SYSTEM_PROMPT_TEMPLATE
+        self.additional_instructions = additional_instructions
 
         # Either client or connector must be provided
         if not client and not connector:
@@ -55,6 +75,7 @@ class MCPAgent:
 
         self._agent: LangChainAgent | None = None
         self._session: MCPSession | None = None
+        self._system_message: SystemMessage | None = None
 
     async def initialize(self) -> None:
         """Initialize the MCP client and agent."""
@@ -71,38 +92,104 @@ class MCPAgent:
             await connector_to_use.connect()
             await connector_to_use.initialize()
 
+        # Create the system message based on available tools
+        await self._create_system_message(connector_to_use)
+
         # Create the agent
         self._agent = LangChainAgent(
-            connector=connector_to_use, llm=self.llm, max_steps=self.max_steps
+            connector=connector_to_use,
+            llm=self.llm,
+            max_steps=self.max_steps,
+            system_message=(self._system_message.content if self._system_message else None),
         )
 
         # Initialize the agent
         await self._agent.initialize()
         self._initialized = True
 
-    async def close(self) -> None:
-        """Close the MCP connection."""
-        try:
-            if self._agent:
-                # Clean up the agent first
-                self._agent = None
+    async def _create_system_message(self, connector: BaseConnector) -> None:
+        """Create the system message based on available tools.
 
-            # If using client with session, close the session through client
-            if self.client and self._session:
-                await self.client.close_session(self.server_name)
-            # If using direct connector, disconnect
-            elif self.connector:
-                await self.connector.disconnect()
+        Args:
+            connector: The connector with available tools.
+        """
+        # If a complete system prompt was provided, use it directly
+        if self.system_prompt:
+            self._system_message = SystemMessage(content=self.system_prompt)
+            return
 
-            self._initialized = False
-        except Exception as e:
-            logger.warning(f"Warning: Error during agent closure: {e}")
-            # Still try to clean up even if there was an error
-            self._agent = None
-            self._initialized = False
+        # Otherwise, build the system prompt from the template and tool descriptions
+        tools = connector.tools
+
+        # Generate tool descriptions
+        tool_descriptions = []
+        for tool in tools:
+            description = f"- {tool.name}: {tool.description}"
+            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)
+        )
+
+        # Add any additional instructions
+        if self.additional_instructions:
+            system_prompt += f"\n\n{self.additional_instructions}"
+
+        # Create the system message
+        self._system_message = SystemMessage(content=system_prompt)
+        logger.info(f"Created system message with {len(tools)} tool descriptions")
+
+    def get_conversation_history(self) -> list[BaseMessage]:
+        """Get the current conversation history.
+
+        Returns:
+            The list of conversation messages.
+        """
+        return self._conversation_history
+
+    def clear_conversation_history(self) -> None:
+        """Clear the conversation history."""
+        self._conversation_history = []
+
+        # Re-add the system message if it exists
+        if self._system_message:
+            self._conversation_history = [self._system_message]
+
+    def add_to_history(self, message: BaseMessage) -> None:
+        """Add a message to the conversation history.
+
+        Args:
+            message: The message to add.
+        """
+        self._conversation_history.append(message)
+
+    def get_system_message(self) -> SystemMessage | None:
+        """Get the current system message.
+
+        Returns:
+            The current system message, or None if not set.
+        """
+        return self._system_message
+
+    def set_system_message(self, message: str) -> None:
+        """Set a new system message.
+
+        Args:
+            message: The new system message content.
+        """
+        self._system_message = SystemMessage(content=message)
+
+        # Update the agent if initialized
+        if self._agent:
+            self._agent.set_system_message(message)
 
     async def run(
-        self, query: str, max_steps: int | None = None, manage_connector: bool = True
+        self,
+        query: str,
+        max_steps: int | None = None,
+        manage_connector: bool = True,
+        external_history: list[BaseMessage] | None = None,
     ) -> str:
         """Run a query using the MCP tools.
 
@@ -116,34 +203,106 @@ class MCPAgent:
                 If True, this method will connect, initialize, and disconnect from
                 the connector automatically. If False, the caller is responsible
                 for managing the connector lifecycle.
+            external_history: Optional external history to use instead of the
+                internal conversation history.
 
         Returns:
             The result of running the query.
         """
         result = ""
+        initialized_here = False
+
         try:
-            if manage_connector:
-                # Initialize if needed
-                if not self._initialized or not self._agent:
-                    await self.initialize()
+            # Initialize if needed
+            if manage_connector and (not self._initialized or not self._agent):
+                logger.info("Initializing agent before running query")
+                await self.initialize()
+                initialized_here = True
+            elif not self._initialized and self.auto_initialize:
+                logger.info("Auto-initializing agent before running query")
+                await self.initialize()
+                initialized_here = True
+
+            # Check if initialization succeeded
+            if not self._agent:
+                raise RuntimeError("MCP agent failed to initialize")
+
+            # Add the user query to conversation history if memory is enabled
+            if self.memory_enabled:
+                self.add_to_history(HumanMessage(content=query))
 
-                # Run the query
-                if not self._agent:
-                    raise RuntimeError("MCP client failed to initialize")
+            # Use the provided history or the internal history
+            history_to_use = (
+                external_history if external_history is not None else self._conversation_history
+            )
 
-                result = await self._agent.run(query, max_steps)
-            else:
-                # Caller is managing connector lifecycle
-                if not self._initialized and self.auto_initialize:
-                    await self.initialize()
+            # Convert messages to format expected by LangChain
+            langchain_history = []
+            for msg in history_to_use:
+                if isinstance(msg, HumanMessage):
+                    langchain_history.append({"type": "human", "content": msg.content})
+                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
 
-                if not self._agent:
-                    raise RuntimeError("MCP client is not initialized")
+            # Run the query with the specified max_steps or default
+            logger.info(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,
+            )
 
-                result = await self._agent.run(query, max_steps)
+            # Add the response to conversation history if memory is enabled
+            if self.memory_enabled:
+                self.add_to_history(AIMessage(content=result))
 
             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
+            if initialized_here and manage_connector:
+                logger.info("Cleaning up resources after initialization error")
+                await self.close()
+            raise
+
         finally:
-            # Make sure to clean up the connection if we're managing it
-            if manage_connector and not self.client:
+            # Clean up resources if we're managing the connector and
+            # we're not using a client that manages sessions
+            if manage_connector and not self.client and not initialized_here:
+                logger.info("Closing agent after query completion")
                 await self.close()
+
+    async def close(self) -> None:
+        """Close the MCP connection with improved error handling."""
+        try:
+            if self._agent:
+                # Clean up the agent first
+                logger.debug("Cleaning up agent")
+                self._agent = None
+
+            # If using client with session, close the session through client
+            if self.client and self._session:
+                logger.debug("Closing session through client")
+                await self.client.close_session(self.server_name)
+                self._session = None
+            # If using direct connector, disconnect
+            elif self.connector:
+                logger.debug("Disconnecting connector")
+                await self.connector.disconnect()
+
+            self._initialized = False
+            logger.info("Agent closed successfully")
+
+        except Exception as 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._initialized = False

mcp_use/agents/prompts/default.py

@@ -0,0 +1,11 @@
+DEFAULT_SYSTEM_PROMPT_TEMPLATE = """You are an assistant with access to these tools:
+
+{tool_descriptions}
+
+Proactively use these tools to:
+- Find real-time information (weather, news, prices)
+- Perform web searches and extract relevant data
+- Execute multi-step tasks by combining tools
+
+You CAN access current information using your tools. Never claim you lack access to real-time data.
+"""

mcp_use/client.py

@@ -9,6 +9,7 @@ import json
 from typing import Any
 
 from .config import create_connector_from_config, load_config_file
+from .logging import logger
 from .session import MCPSession
 
 
@@ -183,40 +184,55 @@ class MCPClient:
         Raises:
             ValueError: If no active session exists or the specified session doesn't exist.
         """
-        session = self.get_session(server_name)
-        await session.disconnect()
-
-        # Remove the session
+        # Determine which server to close
         if server_name is None:
+            if self.active_session is None:
+                logger.warning("No active session to close")
+                return
             server_name = self.active_session
 
-        if server_name in self.sessions:
+        # Check if the session exists
+        if server_name not in self.sessions:
+            logger.warning(f"No session exists for server '{server_name}', nothing to close")
+            return
+
+        # Get the session
+        session = self.sessions[server_name]
+
+        try:
+            # Disconnect from the session
+            logger.info(f"Closing session for server '{server_name}'")
+            await session.disconnect()
+        except Exception as e:
+            logger.error(f"Error closing session for server '{server_name}': {e}")
+        finally:
+            # Remove the session regardless of whether disconnect succeeded
             del self.sessions[server_name]
 
-        # If we closed the active session, set active_session to None
-        if server_name == self.active_session:
-            self.active_session = None
+            # If we closed the active session, set active_session to None
+            if server_name == self.active_session:
+                self.active_session = None
 
     async def close_all_sessions(self) -> None:
-        """Close all active sessions."""
-        for server_name in list(self.sessions.keys()):
-            await self.close_session(server_name)
-
-    async def __aenter__(self) -> "MCPClient":
-        """Enter the async context manager.
-
-        Creates a session for the first available server if no sessions exist.
-
-        Returns:
-            The client instance.
-        """
-        if not self.sessions and self.config.get("mcpServers"):
-            await self.create_session()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Exit the async context manager.
+        """Close all active sessions.
 
-        Closes all active sessions.
+        This method ensures all sessions are closed even if some fail.
         """
-        await self.close_all_sessions()
+        # Get a list of all session names first to avoid modification during iteration
+        server_names = list(self.sessions.keys())
+        errors = []
+
+        for server_name in server_names:
+            try:
+                logger.info(f"Closing session for server '{server_name}'")
+                await self.close_session(server_name)
+            except Exception as e:
+                error_msg = f"Failed to close session for server '{server_name}': {e}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+
+        # Log summary if there were errors
+        if errors:
+            logger.error(f"Encountered {len(errors)} errors while closing sessions")
+        else:
+            logger.info("All sessions closed successfully")