letta-nightly 0.6.40.dev20250314173529__py3-none-any.whl → 0.6.40.dev20250314222759__py3-none-any.whl

letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.39"
+__version__ = "0.6.40"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client

letta/agent.py

@@ -22,11 +22,11 @@ from letta.errors import ContextWindowExceededError
 from letta.functions.ast_parsers import coerce_dict_args_by_annotations, get_function_annotations_from_source
 from letta.functions.functions import get_function_from_module
 from letta.functions.helpers import execute_composio_action, generate_composio_action_from_func_name
+from letta.functions.mcp_client.base_client import BaseMCPClient
 from letta.helpers import ToolRulesSolver
 from letta.helpers.composio_helpers import get_composio_api_key
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
-from letta.helpers.mcp_helpers import BaseMCPClient
 from letta.interface import AgentInterface
 from letta.llm_api.helpers import calculate_summarizer_cutoff, get_token_counts_for_messages, is_context_overflow_error
 from letta.llm_api.llm_api_tools import create

letta/functions/mcp_client/base_client.py

@@ -0,0 +1,61 @@
+import asyncio
+from typing import List, Optional, Tuple
+
+from mcp import ClientSession, Tool
+
+from letta.functions.mcp_client.types import BaseServerConfig
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
+
+class BaseMCPClient:
+    def __init__(self):
+        self.session: Optional[ClientSession] = None
+        self.stdio = None
+        self.write = None
+        self.initialized = False
+        self.loop = asyncio.new_event_loop()
+        self.cleanup_funcs = []
+
+    def connect_to_server(self, server_config: BaseServerConfig):
+        asyncio.set_event_loop(self.loop)
+        success = self._initialize_connection(server_config)
+
+        if success:
+            self.loop.run_until_complete(self.session.initialize())
+            self.initialized = True
+        else:
+            raise RuntimeError(
+                f"Connecting to MCP server failed. Please review your server config: {server_config.model_dump_json(indent=4)}"
+            )
+
+    def _initialize_connection(self, server_config: BaseServerConfig) -> bool:
+        raise NotImplementedError("Subclasses must implement _initialize_connection")
+
+    def list_tools(self) -> List[Tool]:
+        self._check_initialized()
+        response = self.loop.run_until_complete(self.session.list_tools())
+        return response.tools
+
+    def execute_tool(self, tool_name: str, tool_args: dict) -> Tuple[str, bool]:
+        self._check_initialized()
+        result = self.loop.run_until_complete(self.session.call_tool(tool_name, tool_args))
+        return str(result.content), result.isError
+
+    def _check_initialized(self):
+        if not self.initialized:
+            logger.error("MCPClient has not been initialized")
+            raise RuntimeError("MCPClient has not been initialized")
+
+    def cleanup(self):
+        try:
+            for cleanup_func in self.cleanup_funcs:
+                cleanup_func()
+            self.initialized = False
+            if not self.loop.is_closed():
+                self.loop.close()
+        except Exception as e:
+            logger.warning(e)
+        finally:
+            logger.info("Cleaned up MCP clients on shutdown.")

letta/functions/mcp_client/sse_client.py

@@ -0,0 +1,21 @@
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+
+from letta.functions.mcp_client.base_client import BaseMCPClient
+from letta.functions.mcp_client.types import SSEServerConfig
+
+# see: https://modelcontextprotocol.io/quickstart/user
+MCP_CONFIG_TOPLEVEL_KEY = "mcpServers"
+
+
+class SSEMCPClient(BaseMCPClient):
+    def _initialize_connection(self, server_config: SSEServerConfig) -> bool:
+        sse_cm = sse_client(url=server_config.server_url)
+        sse_transport = self.loop.run_until_complete(sse_cm.__aenter__())
+        self.stdio, self.write = sse_transport
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(sse_cm.__aexit__(None, None, None)))
+
+        session_cm = ClientSession(self.stdio, self.write)
+        self.session = self.loop.run_until_complete(session_cm.__aenter__())
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
+        return True

letta/functions/mcp_client/stdio_client.py

@@ -0,0 +1,103 @@
+import sys
+from contextlib import asynccontextmanager
+
+import anyio
+import anyio.lowlevel
+import mcp.types as types
+from anyio.streams.text import TextReceiveStream
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import get_default_environment
+
+from letta.functions.mcp_client.base_client import BaseMCPClient
+from letta.functions.mcp_client.types import StdioServerConfig
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
+
+class StdioMCPClient(BaseMCPClient):
+    def _initialize_connection(self, server_config: StdioServerConfig) -> bool:
+        try:
+            server_params = StdioServerParameters(command=server_config.command, args=server_config.args)
+            stdio_cm = forked_stdio_client(server_params)
+            stdio_transport = self.loop.run_until_complete(stdio_cm.__aenter__())
+            self.stdio, self.write = stdio_transport
+            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(stdio_cm.__aexit__(None, None, None)))
+
+            session_cm = ClientSession(self.stdio, self.write)
+            self.session = self.loop.run_until_complete(session_cm.__aenter__())
+            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
+            return True
+        except Exception:
+            return False
+
+
+@asynccontextmanager
+async def forked_stdio_client(server: StdioServerParameters):
+    """
+    Client transport for stdio: this will connect to a server by spawning a
+    process and communicating with it over stdin/stdout.
+    """
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    try:
+        process = await anyio.open_process(
+            [server.command, *server.args],
+            env=server.env or get_default_environment(),
+            stderr=sys.stderr,  # Consider logging stderr somewhere instead of silencing it
+        )
+    except OSError as exc:
+        raise RuntimeError(f"Failed to spawn process: {server.command} {server.args}") from exc
+
+    async def stdout_reader():
+        assert process.stdout, "Opened process is missing stdout"
+        buffer = ""
+        try:
+            async with read_stream_writer:
+                async for chunk in TextReceiveStream(
+                    process.stdout,
+                    encoding=server.encoding,
+                    errors=server.encoding_error_handler,
+                ):
+                    lines = (buffer + chunk).split("\n")
+                    buffer = lines.pop()
+                    for line in lines:
+                        try:
+                            message = types.JSONRPCMessage.model_validate_json(line)
+                        except Exception as exc:
+                            await read_stream_writer.send(exc)
+                            continue
+                        await read_stream_writer.send(message)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stdin_writer():
+        assert process.stdin, "Opened process is missing stdin"
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    json = message.model_dump_json(by_alias=True, exclude_none=True)
+                    await process.stdin.send(
+                        (json + "\n").encode(
+                            encoding=server.encoding,
+                            errors=server.encoding_error_handler,
+                        )
+                    )
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def watch_process_exit():
+        returncode = await process.wait()
+        if returncode != 0:
+            raise RuntimeError(f"Subprocess exited with code {returncode}. Command: {server.command} {server.args}")
+
+    async with anyio.create_task_group() as tg, process:
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+        tg.start_soon(watch_process_exit)
+
+        with anyio.move_on_after(0.2):
+            await anyio.sleep_forever()
+
+        yield read_stream, write_stream

letta/functions/mcp_client/types.py

@@ -0,0 +1,48 @@
+from enum import Enum
+from typing import List, Optional
+
+from mcp import Tool
+from pydantic import BaseModel, Field
+
+
+class MCPTool(Tool):
+    """A simple wrapper around MCP's tool definition (to avoid conflict with our own)"""
+
+
+class MCPServerType(str, Enum):
+    SSE = "sse"
+    STDIO = "stdio"
+
+
+class BaseServerConfig(BaseModel):
+    server_name: str = Field(..., description="The name of the server")
+    type: MCPServerType
+
+
+class SSEServerConfig(BaseServerConfig):
+    type: MCPServerType = MCPServerType.SSE
+    server_url: str = Field(..., description="The URL of the server (MCP SSE client will connect to this URL)")
+
+    def to_dict(self) -> dict:
+        values = {
+            "transport": "sse",
+            "url": self.server_url,
+        }
+        return values
+
+
+class StdioServerConfig(BaseServerConfig):
+    type: MCPServerType = MCPServerType.STDIO
+    command: str = Field(..., description="The command to run (MCP 'local' client will run this command)")
+    args: List[str] = Field(..., description="The arguments to pass to the command")
+    env: Optional[dict[str, str]] = Field(None, description="Environment variables to set")
+
+    def to_dict(self) -> dict:
+        values = {
+            "transport": "stdio",
+            "command": self.command,
+            "args": self.args,
+        }
+        if self.env is not None:
+            values["env"] = self.env
+        return values

letta/functions/schema_generator.py

@@ -6,7 +6,7 @@ from composio.client.collections import ActionParametersModel
 from docstring_parser import parse
 from pydantic import BaseModel
 
-from letta.helpers.mcp_helpers import MCPTool
+from letta.functions.mcp_client.types import MCPTool
 
 
 def is_optional(annotation):

letta/schemas/tool.py

@@ -17,12 +17,12 @@ from letta.functions.helpers import (
     generate_mcp_tool_wrapper,
     generate_model_from_args_json_schema,
 )
+from letta.functions.mcp_client.types import MCPTool
 from letta.functions.schema_generator import (
     generate_schema_from_args_schema_v2,
     generate_tool_schema_for_composio,
     generate_tool_schema_for_mcp,
 )
-from letta.helpers.mcp_helpers import MCPTool
 from letta.log import get_logger
 from letta.orm.enums import ToolType
 from letta.schemas.letta_base import LettaBase

letta/server/rest_api/routers/v1/tools.py

@@ -12,8 +12,8 @@ from composio.exceptions import (
 from fastapi import APIRouter, Body, Depends, Header, HTTPException
 
 from letta.errors import LettaToolCreateError
+from letta.functions.mcp_client.types import MCPTool, SSEServerConfig, StdioServerConfig
 from letta.helpers.composio_helpers import get_composio_api_key
-from letta.helpers.mcp_helpers import MCPTool, SSEServerConfig, StdioServerConfig
 from letta.log import get_logger
 from letta.orm.errors import UniqueConstraintViolationError
 from letta.schemas.letta_message import ToolReturnMessage

letta/server/server.py

@@ -20,18 +20,12 @@ from letta.agent import Agent, save_agent
 from letta.config import LettaConfig
 from letta.data_sources.connectors import DataConnector, load_data
 from letta.dynamic_multi_agent import DynamicMultiAgent
+from letta.functions.mcp_client.base_client import BaseMCPClient
+from letta.functions.mcp_client.sse_client import MCP_CONFIG_TOPLEVEL_KEY, SSEMCPClient
+from letta.functions.mcp_client.stdio_client import StdioMCPClient
+from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerConfig, StdioServerConfig
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
-from letta.helpers.mcp_helpers import (
-    MCP_CONFIG_TOPLEVEL_KEY,
-    BaseMCPClient,
-    MCPServerType,
-    MCPTool,
-    SSEMCPClient,
-    SSEServerConfig,
-    StdioMCPClient,
-    StdioServerConfig,
-)
 
 # TODO use custom interface
 from letta.interface import AgentInterface  # abstract
@@ -343,11 +337,12 @@ class SyncServer(Server):
                 self.mcp_clients[server_name] = StdioMCPClient()
             else:
                 raise ValueError(f"Invalid MCP server config: {server_config}")
+
             try:
                 self.mcp_clients[server_name].connect_to_server(server_config)
-            except:
-                logger.exception(f"Failed to connect to MCP server: {server_name}")
-                raise
+            except Exception as e:
+                logger.error(e)
+                self.mcp_clients.pop(server_name)
 
         # Print out the tools that are connected
         for server_name, client in self.mcp_clients.items():

{letta_nightly-0.6.40.dev20250314173529.dist-info → letta_nightly-0.6.40.dev20250314222759.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.6.40.dev20250314173529
+Version: 0.6.40.dev20250314222759
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team

{letta_nightly-0.6.40.dev20250314173529.dist-info → letta_nightly-0.6.40.dev20250314222759.dist-info}/RECORD

@@ -1,6 +1,6 @@
-letta/__init__.py,sha256=q6k2pwvx7XD4fcSf0faPYn9dD9Yj6778hhgrULdXd6s,918
+letta/__init__.py,sha256=UHqXlWcxu1BBS0f_6UszKphaK2mFq0xvQby9G3WcIFc,918
 letta/__main__.py,sha256=6Hs2PV7EYc5Tid4g4OtcLXhqVHiNYTGzSBdoOnW2HXA,29
-letta/agent.py,sha256=CK6E-ZeFOuzoJxgegogDNVcgfUMS9SAebpYI7HTvS2Q,68128
+letta/agent.py,sha256=xBXIgaeiGXDBpgNKX0-2w10iRBBvcLgBK0W85jTiSX8,68141
 letta/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 letta/agents/base_agent.py,sha256=8IMB7UK4ft-Wi-ZYjX7akqQUhk_cSswRgepqeZyvCMs,1550
 letta/agents/ephemeral_agent.py,sha256=DBMXT50UQoqjkvl_Piwle3Fy7iXopy15oMWwnWzbpvo,2751
@@ -29,13 +29,17 @@ letta/functions/function_sets/multi_agent.py,sha256=QTRvK7G-B_2ShIP7hw6Xw50P5VMM
 letta/functions/functions.py,sha256=NyWLh7a-f4mXti5vM1oWDwXzfA58VmVVqL03O9vosKY,5672
 letta/functions/helpers.py,sha256=0I-ezZeM3slhAifpdlR5k2Rn6GxExD6xACCKuoYmE8M,29119
 letta/functions/interface.py,sha256=s_PPp5WDvGH_y9KUpMlORkdC141ITczFk3wsevrrUD8,2866
-letta/functions/schema_generator.py,sha256=wbV8tDx5RIkrSP9JNVo3PUuulN-WZie1WWLg69QTOjE,22320
+letta/functions/mcp_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+letta/functions/mcp_client/base_client.py,sha256=uHp4XuXFy3avCkc3XYr1IxEPTHZWADqrDk-kR-oFyLY,2144
+letta/functions/mcp_client/sse_client.py,sha256=VElaM87sEMTvam4t3IJgWGhyqs5BL4YVErqVPv2hQug,961
+letta/functions/mcp_client/stdio_client.py,sha256=qy9vHB99oI4pCAQ4KmQUKY0bTboaQI4pnPm9DPQfPvY,4210
+letta/functions/mcp_client/types.py,sha256=nmcnQn2EpxXzXg5_pWPsHZobfxO6OucaUgz1bVvam7o,1411
+letta/functions/schema_generator.py,sha256=4hiDQpHemyfKWME-5X6xJuSiv7g9V_BgAPFegohHBIM,22327
 letta/helpers/__init__.py,sha256=p0luQ1Oe3Skc6sH4O58aHHA3Qbkyjifpuq0DZ1GAY0U,59
 letta/helpers/composio_helpers.py,sha256=6CWV483vE3N-keQlblexxBiQHxorMAgQuvbok4adGO4,949
 letta/helpers/converters.py,sha256=3qHoPDPa7ycPeeE6eLYZ8mad1mA6oGPgY95lDyhb3_A,8971
 letta/helpers/datetime_helpers.py,sha256=7U5ZJkE0cLki4sG8ukIHZSAoFfQpLWQu2kFekETy6Zg,2633
 letta/helpers/json_helpers.py,sha256=PWZ5HhSqGXO4e563dM_8M72q7ScirjXQ4Rv1ckohaV8,396
-letta/helpers/mcp_helpers.py,sha256=T1Jz2emYlDBbTsqZKckUGWiDjeoD_7lOUGx7GWU1mM8,4789
 letta/helpers/tool_execution_helper.py,sha256=bskCscuz2nqoUboFcYA7sQGeikdEyqiYnNpO4gLQTdc,7469
 letta/helpers/tool_rule_solver.py,sha256=z-2Zq_qWykgWanFZYxtxUee4FkMnxqvntXe2tomoH68,6774
 letta/humans/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -199,7 +203,7 @@ letta/schemas/run.py,sha256=SRqPRziINIiPunjOhE_NlbnQYgxTvqmbauni_yfBQRA,2085
 letta/schemas/sandbox_config.py,sha256=SZCo3FSMz-DIBMKGu0atT4tsVFXGsqMFPaJnjrxpkX4,5993
 letta/schemas/source.py,sha256=-BQVolcXA2ziCu2ztR6cbTdGUc8G7vGJy7rvpdf1hpg,2880
 letta/schemas/step.py,sha256=WkcVnruUUOWLKwiWPn2Gfal4EQZPNLqlsd9859xhgsw,2224
-letta/schemas/tool.py,sha256=A6HS_nHoxAOVyv4u6NMxzMNsp0PG40kY2uDS8FA06M4,12284
+letta/schemas/tool.py,sha256=PXWxEqzM-kADijlsJzu0ZYtWLnjpq4ZUNX4NzesIWNQ,12291
 letta/schemas/tool_rule.py,sha256=2YQZba4fXS3u4j8pIk7BDujfq8rnxSVMwJSyaVgApH4,2149
 letta/schemas/usage.py,sha256=8oYRH-JX0PfjIu2zkT5Uu3UWQ7Unnz_uHiO8hRGI4m0,912
 letta/schemas/user.py,sha256=V32Tgl6oqB3KznkxUz12y7agkQicjzW7VocSpj78i6Q,1526
@@ -243,12 +247,12 @@ letta/server/rest_api/routers/v1/sandbox_configs.py,sha256=9hqnnMwJ3wCwO-Bezu3Xl
 letta/server/rest_api/routers/v1/sources.py,sha256=rpQhaRHyzGUK43LX623L8BBLqL85HJ6fUYPMvI4k3kA,8434
 letta/server/rest_api/routers/v1/steps.py,sha256=DVVwaxLNbNAgWpr2oQkrNjdS-wi0bP8kVJZUO-hiaf8,3275
 letta/server/rest_api/routers/v1/tags.py,sha256=coydgvL6-9cuG2Hy5Ea7QY3inhTHlsf69w0tcZenBus,880
-letta/server/rest_api/routers/v1/tools.py,sha256=UVEl7nQr74C-VV02Enw6dx9S8UvPMnM97Ju-X92GswA,17204
+letta/server/rest_api/routers/v1/tools.py,sha256=pWIlYUksq9QNFSmknB_DdSB5zs77zaHoj-Ha-WxYkO8,17211
 letta/server/rest_api/routers/v1/users.py,sha256=G5DBHSkPfBgVHN2Wkm-rVYiLQAudwQczIq2Z3YLdbVo,2277
 letta/server/rest_api/routers/v1/voice.py,sha256=7J0L-Nkz65m0PXcpQI2ATMIZzumDDSUzgtIus7d-tv8,2461
 letta/server/rest_api/static_files.py,sha256=NG8sN4Z5EJ8JVQdj19tkFa9iQ1kBPTab9f_CUxd_u4Q,3143
 letta/server/rest_api/utils.py,sha256=aF0u__Q33-aPWAiHi9JA0jKAjqnwbVKzJdD5NgFpnOU,13828
-letta/server/server.py,sha256=w2vCdTz4rETNA_pmVmEMniUIygIlIeD6Qud_dZNcfVQ,75772
+letta/server/server.py,sha256=Ml4-MklB28x3iAONfp3m2oSbrhVctAvvxUWApeWIWrk,75884
 letta/server/startup.sh,sha256=eo7zz4HGu5ryOshfbOSGbXpUDDyoaP7fTq4z8269uaw,1939
 letta/server/static_files/assets/index-048c9598.js,sha256=mR16XppvselwKCcNgONs4L7kZEVa4OEERm4lNZYtLSk,146819
 letta/server/static_files/assets/index-0e31b727.css,sha256=SBbja96uiQVLDhDOroHgM6NSl7tS4lpJRCREgSS_hA8,7672
@@ -290,8 +294,8 @@ letta/supervisor_multi_agent.py,sha256=dw6XPAxZcyjAXEVYkMIxJFhZdR2m2_rq-fYev5hZE
 letta/system.py,sha256=dnOrS2FlRMwijQnOvfrky0Lg8wEw-FUq2zzfAJOUSKA,8477
 letta/tracing.py,sha256=h_-c2lIKHmA7yCLOvgaHijMabmRC__FAl2rZtVKufUM,8017
 letta/utils.py,sha256=AdHrQ2OQ3V4XhJ1LtYwbLUO71j2IJY37cIUxXPgaaRY,32125
-letta_nightly-0.6.40.dev20250314173529.dist-info/LICENSE,sha256=mExtuZ_GYJgDEI38GWdiEYZizZS4KkVt2SF1g_GPNhI,10759
-letta_nightly-0.6.40.dev20250314173529.dist-info/METADATA,sha256=vByuTR-Z0aV0d2OsTuIqBwOJjtb41l7WPqM62GTdUvQ,22886
-letta_nightly-0.6.40.dev20250314173529.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-letta_nightly-0.6.40.dev20250314173529.dist-info/entry_points.txt,sha256=2zdiyGNEZGV5oYBuS-y2nAAgjDgcC9yM_mHJBFSRt5U,40
-letta_nightly-0.6.40.dev20250314173529.dist-info/RECORD,,
+letta_nightly-0.6.40.dev20250314222759.dist-info/LICENSE,sha256=mExtuZ_GYJgDEI38GWdiEYZizZS4KkVt2SF1g_GPNhI,10759
+letta_nightly-0.6.40.dev20250314222759.dist-info/METADATA,sha256=_O99vV3C3JQYXICthAEXivxmz4g0rYe2ByP1bNOXu6Y,22886
+letta_nightly-0.6.40.dev20250314222759.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+letta_nightly-0.6.40.dev20250314222759.dist-info/entry_points.txt,sha256=2zdiyGNEZGV5oYBuS-y2nAAgjDgcC9yM_mHJBFSRt5U,40
+letta_nightly-0.6.40.dev20250314222759.dist-info/RECORD,,

letta/helpers/mcp_helpers.py

@@ -1,129 +0,0 @@
-import asyncio
-from enum import Enum
-from typing import List, Optional, Tuple
-
-from mcp import ClientSession, StdioServerParameters, Tool
-from mcp.client.sse import sse_client
-from mcp.client.stdio import stdio_client
-from pydantic import BaseModel, Field
-
-from letta.log import get_logger
-
-logger = get_logger(__name__)
-
-# see: https://modelcontextprotocol.io/quickstart/user
-MCP_CONFIG_TOPLEVEL_KEY = "mcpServers"
-
-
-class MCPTool(Tool):
-    """A simple wrapper around MCP's tool definition (to avoid conflict with our own)"""
-
-
-class MCPServerType(str, Enum):
-    SSE = "sse"
-    STDIO = "stdio"
-
-
-class BaseServerConfig(BaseModel):
-    server_name: str = Field(..., description="The name of the server")
-    type: MCPServerType
-
-
-class SSEServerConfig(BaseServerConfig):
-    type: MCPServerType = MCPServerType.SSE
-    server_url: str = Field(..., description="The URL of the server (MCP SSE client will connect to this URL)")
-
-    def to_dict(self) -> dict:
-        values = {
-            "transport": "sse",
-            "url": self.server_url,
-        }
-        return values
-
-
-class StdioServerConfig(BaseServerConfig):
-    type: MCPServerType = MCPServerType.STDIO
-    command: str = Field(..., description="The command to run (MCP 'local' client will run this command)")
-    args: List[str] = Field(..., description="The arguments to pass to the command")
-    env: Optional[dict[str, str]] = Field(None, description="Environment variables to set")
-
-    def to_dict(self) -> dict:
-        values = {
-            "transport": "stdio",
-            "command": self.command,
-            "args": self.args,
-        }
-        if self.env is not None:
-            values["env"] = self.env
-        return values
-
-
-class BaseMCPClient:
-    def __init__(self):
-        self.session: Optional[ClientSession] = None
-        self.stdio = None
-        self.write = None
-        self.initialized = False
-        self.loop = asyncio.new_event_loop()
-        self.cleanup_funcs = []
-
-    def connect_to_server(self, server_config: BaseServerConfig):
-        asyncio.set_event_loop(self.loop)
-        self._initialize_connection(server_config)
-        self.loop.run_until_complete(self.session.initialize())
-        self.initialized = True
-
-    def _initialize_connection(self, server_config: BaseServerConfig):
-        raise NotImplementedError("Subclasses must implement _initialize_connection")
-
-    def list_tools(self) -> List[Tool]:
-        self._check_initialized()
-        response = self.loop.run_until_complete(self.session.list_tools())
-        return response.tools
-
-    def execute_tool(self, tool_name: str, tool_args: dict) -> Tuple[str, bool]:
-        self._check_initialized()
-        result = self.loop.run_until_complete(self.session.call_tool(tool_name, tool_args))
-        return str(result.content), result.isError
-
-    def _check_initialized(self):
-        if not self.initialized:
-            logger.error("MCPClient has not been initialized")
-            raise RuntimeError("MCPClient has not been initialized")
-
-    def cleanup(self):
-        try:
-            for cleanup_func in self.cleanup_funcs:
-                cleanup_func()
-            self.initialized = False
-            if not self.loop.is_closed():
-                self.loop.close()
-        except Exception as e:
-            logger.warning(e)
-        finally:
-            logger.info("Cleaned up MCP clients on shutdown.")
-
-
-class StdioMCPClient(BaseMCPClient):
-    def _initialize_connection(self, server_config: StdioServerConfig):
-        server_params = StdioServerParameters(command=server_config.command, args=server_config.args)
-        stdio_cm = stdio_client(server_params)
-        stdio_transport = self.loop.run_until_complete(stdio_cm.__aenter__())
-        self.stdio, self.write = stdio_transport
-        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(stdio_cm.__aexit__(None, None, None)))
-
-        session_cm = ClientSession(self.stdio, self.write)
-        self.session = self.loop.run_until_complete(session_cm.__aenter__())
-        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
-
-
-class SSEMCPClient(BaseMCPClient):
-    def _initialize_connection(self, server_config: SSEServerConfig):
-        sse_cm = sse_client(url=server_config.server_url)
-        sse_transport = self.loop.run_until_complete(sse_cm.__aenter__())
-        self.stdio, self.write = sse_transport
-        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(sse_cm.__aexit__(None, None, None)))
-
-        session_cm = ClientSession(self.stdio, self.write)
-        self.session = self.loop.run_until_complete(session_cm.__aenter__())
-        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))