feat: 支持 SSE 连接

Author: wangtao2001

docs/tutorials/agent/agent.md

@@ -309,17 +309,15 @@ import asyncio
 qwen = Qwen()
 
 async def main():
-    async with MCPServer(
-        type='stdio',
-        command='uv',
-        args=['--directory', 'examples/agent', 'run', 'mcp_server.py'],
-    ) as mcp_server:
+    async with MCPServer( {
+        'command': 'uv',
+        'args': ['--directory', 'examples/agent', 'run', 'mcp_server.py'],
+    } ) as mcp_server:
 
         agent = Agent(
             llm=qwen,
             mcp_server=[mcp_server]
         )
-        await agent.initialize()
         controller = Controller()
         _, resp = await controller.run(agent, "帮我查询南京今天的天气")
         print(resp)
@@ -328,12 +326,10 @@ if __name__ == '__main__':
     asyncio.run(main())
 ```
 
-这里有三需要注意:
+这里有两点需要注意:
 
 - MCP Server 必须使用 `async with` 语句块来启动。
 
-- 如果给 `Agent` 传递了 `mcp_server` 参数, 必须调用 `await agent.initialize()` 方法等待初始化完成。
-
 - `Controller` 必须使用异步的 `run` 方法来启动。不能使用同步的 `run_sync` 方法, 因为本质上 `run_sync` 只是 `asyncio.run(controller.run(...))` 的包装。
 
 ### MCP Server 与外部工具的比较

examples/agent/use_mcp.py

@@ -4,24 +4,24 @@
 # File Name: examples/agent/use_mcp.py
 # Description: 使用 MCP 工具
 
-from course_graph.agent import Agent, Controller, MCPServer, TraceEvent
+from course_graph.agent import Agent, Controller, MCPServer
 from course_graph.llm import Qwen
 import asyncio
 qwen = Qwen()
 
 
 async def main():
     async with MCPServer(
-        type='stdio',
-        command='uv',
-        args=['--directory', 'examples/agent', 'run', 'mcp_server.py'],
+        {
+            'command': 'uv',
+            'args': ['--directory', 'examples/agent', 'run', 'mcp_server.py'],
+        }
     ) as mcp_server:
 
         agent = Agent(
             llm=qwen,
             mcp_server=[mcp_server]
         )
-        await agent.initialize()
         controller = Controller()
         _, resp = await controller.run(agent, "帮我查询今天南京的天气")
         print(resp)

src/course_graph/agent/__init__.py

@@ -8,4 +8,4 @@
 from .agent import Agent
 from .controller import Controller
 from .types import Result, ContextVariables, TraceEvent
-from .mcp import MCPServer
+from .mcp import MCPServer, STDIO, SSE

src/course_graph/agent/agent.py

@@ -69,13 +69,9 @@ def __init__(
             self.tool_choice = tool_choice
 
         self.messages: list[ChatCompletionMessageParam] = []
-        self.mcp_server = mcp_server
 
-    async def initialize(self):
-        """ 等待初始化智能体
-        """
-        for server in self.mcp_server:
-            tools = await server.list_tools()
+        for server in mcp_server:
+            tools = server.tools
             for tool in tools:
                 self.tools.append({
                     'type': 'function',

src/course_graph/agent/mcp.py

@@ -4,36 +4,48 @@
 # File Name: course_graph/agent/mcp.py
 # Description: MCP Server 和 Client 实现相关
 
-from typing import Literal
+from typing import TypedDict, Required, NotRequired
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
+from mcp.client.sse import sse_client
 from contextlib import AsyncExitStack
 from mcp.types import Tool
 
 
+class STDIO(TypedDict):
+    command: Required[str]
+    args: Required[list[str]]
+    envs: NotRequired[dict[str, str]]
+
+
+class SSE(TypedDict):
+    url: Required[str]
+    headers: NotRequired[dict[str, str]]
+
+
 class MCPServer:
-    def __init__(self, type: Literal['stdio'], command: str, args: list[str], envs: dict[str, str] = None):
+    def __init__(self, server: STDIO | SSE):
         """ MCP 服务器
 
         Args:
-            type (Literal['stdio']): 服务器类型, 目前只支持 'stdio'
-            command (str): 命令
-            args (list[str]): 参数
-            envs (dict[str, str], optional): 环境变量. Defaults to None.
+            server (STDIO | SSE): MCP 服务器配置
         """
-        self.params = StdioServerParameters(command=command, args=args, envs=envs)
+        self.server = server
         self.stack = AsyncExitStack()
-        self.session = None
+        self.session: ClientSession = None
+        self.tools: list[Tool] = None
 
     async def __aenter__(self):
-        stdio_transport = await self.stack.enter_async_context(stdio_client(self.params))
-        self.stdio, self.write = stdio_transport
-        self.session = await self.stack.enter_async_context(ClientSession(self.stdio, self.write))
+        if 'command' in self.server.keys():
+            self.params = StdioServerParameters(command=self.server['command'], args=self.server['args'], envs=self.server.get('envs'))
+            transport = await self.stack.enter_async_context(stdio_client(self.params))
+        else:
+            transport = await self.stack.enter_async_context(sse_client(url=self.server['url'], headers=self.server.get('headers')))
+        self.session = await self.stack.enter_async_context(ClientSession(transport[0], transport[1]))
         await self.session.initialize()
+        
+        self.tools = (await self.session.list_tools()).tools
         return self
 
-    async def list_tools(self) -> list[Tool]:
-        return (await self.session.list_tools()).tools
-    
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.stack.aclose()