improve ux

Author: vbarda

README.md

@@ -46,7 +46,7 @@ def multiply(a: int, b: int) -> int:
     return a * b
 
 if __name__ == "__main__":
-    mcp.run()
+    mcp.run(transport="stdio")
 ```
 
 ### Client
@@ -103,7 +103,11 @@ async def get_weather(location: str) -> int:
     return "It's always sunny in New York"
 
 if __name__ == "__main__":
-    mcp.run()
+    mcp.run(transport="sse")
+```
+
+```bash
+python weather_server.py
 ```
 
 ### Client
@@ -115,19 +119,21 @@ from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
 model = ChatOpenAI(model="gpt-4o")
 
-async with MultiServerMCPClient() as client:
-    await client.connect_to_server(
-        "math",
-        command="python",
-        # Make sure to update to the full absolute path to your math_server.py file
-        args=["/path/to/math_server.py"],
-    )
-    await client.connect_to_server(
-        "weather",
-        command="python",
-        # Make sure to update to the full absolute path to your weather_server.py file
-        args=["/path/to/weather_server.py"],
-    )
+async with MultiServerMCPClient(
+    {
+        "math": {
+            "command": "python",
+            # Make sure to update to the full absolute path to your math_server.py file
+            "args": ["/path/to/math_server.py"],
+            "transport": "stdio",
+        },
+        "weather": {
+            # make sure you start your weather server on port 8000
+            "url": "http://localhost:8000/sse",
+            "transport": "sse",
+        }
+    }
+) as client:
     agent = create_react_agent(model, client.get_tools())
     math_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
     weather_response = await agent.ainvoke({"messages": "what is the weather in nyc?"})

langchain_mcp_adapters/client.py

@@ -1,6 +1,6 @@
 from contextlib import AsyncExitStack
 from types import TracebackType
-from typing import Literal, cast
+from typing import Literal, TypedDict, cast
 
 from langchain_core.tools import BaseTool
 from mcp import ClientSession, StdioServerParameters
@@ -13,10 +13,25 @@
 DEFAULT_ENCODING_ERROR_HANDLER = "strict"
 
 
+class StdioConnection(TypedDict):
+    transport: Literal["stdio"]
+    command: str
+    args: list[str]
+    env: dict[str, str] | None
+    encoding: str
+    encoding_error_handler: Literal["strict", "ignore", "replace"]
+
+
+class SSEConnection(TypedDict):
+    transport: Literal["sse"]
+    url: str
+
+
 class MultiServerMCPClient:
     """Client for connecting to multiple MCP servers and loading LangChain-compatible tools from them."""
 
-    def __init__(self) -> None:
+    def __init__(self, connections: dict[str, StdioConnection | SSEConnection] = None) -> None:
+        self.connections = connections
         self.exit_stack = AsyncExitStack()
         self.sessions: dict[str, ClientSession] = {}
         self.server_name_to_tools: dict[str, list[BaseTool]] = {}
@@ -151,6 +166,14 @@ def get_tools(self) -> list[BaseTool]:
         return all_tools
 
     async def __aenter__(self) -> "MultiServerMCPClient":
+        connections = self.connections or {}
+        for server_name, connection in connections.items():
+            connection_dict = connection.copy()
+            transport = connection_dict.pop("transport")
+            if transport == "stdio":
+                await self.connect_to_server_via_stdio(server_name, **connection_dict)
+            elif transport == "sse":
+                await self.connect_to_server_via_sse(server_name, **connection_dict)
         return self
 
     async def __aexit__(