feat: MCP Server

Author: wangtao2001

docs/tutorials/agent/agent.md

@@ -75,7 +75,7 @@ assistant.add_too_functions(get_weather)
 当然，在函数定义时就可以声明其为一个工具函数:
 
 ```python
-@assistant.tool
+@assistant.tool()
 def get_weather(location):
     ...
 ```

src/course_graph/agent/__init__.py

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

src/course_graph/agent/agent.py

@@ -13,34 +13,44 @@
 from typing import Callable
 from typing import Literal
 from openai import NOT_GIVEN, NotGiven
+from .mcp import MCPServer
+import asyncio
 
 
 class Agent:
 
     def __init__(
             self,
-            name: str,
             llm: LLMBase,
+            name: str = 'Assistant',
             functions: list[Callable] = None,
             tool_choice: str | NotGiven | Literal['required', 'auto', 'none'] = NOT_GIVEN,
             parallel_tool_calls: bool | NotGiven = NOT_GIVEN,
-            instruction: str | Callable[[ContextVariables], str] | Callable[[], str] = 'You are a helpful assistant.'
+            instruction: str | Callable[[ContextVariables], str] | Callable[[], str] = 'You are a helpful assistant.',
+            mcp_server: list[MCPServer] = None,
+            mcp_impl: Literal['function_call'] = 'function_call'
     ) -> None:
         """ 智能体类
 
-        Args: name (str): 名称 
-        llm (LLMBase): 大模型 
-        functions: (list[Callable], optional): 工具函数. Defaults to None.
-        parallel_tool_calls: (bool, optional): 允许工具并行调用. Defaults to False.
-        tool_choice: (Literal['required', 'auto', 'none'] | NotGiven, optional). 强制使用工具函数, 选择模式或提供函数名称. Defaults to NOT_GIVEN.
-        instruction (str | Callable[[ContextVariables], str] | Callable[[], str], optional): 指令. Defaults to 'You are a helpful assistant.'.
+        Args: 
+            llm (LLMBase): 大模型 
+            name (str, optional): 名称. Defaults to 'Assistant'.
+            functions: (list[Callable], optional): 工具函数. Defaults to None.
+            parallel_tool_calls: (bool, optional): 允许工具并行调用. Defaults to False.
+            tool_choice: (Literal['required', 'auto', 'none'] | NotGiven, optional). 强制使用工具函数, 选择模式或提供函数名称. Defaults to NOT_GIVEN.
+            instruction (str | Callable[[ContextVariables], str] | Callable[[], str], optional): 指令. Defaults to 'You are a helpful assistant.'.
+            mcp_server: (list[MCPServer], optional): MCP 服务器. Defaults to None.
+            mcp_impl: (Literal['function_call'] | NotGiven, optional): MCP 协议实现方式, 目前只支持 'function_call'. Defaults to 'function_call'.
         """
-        self.name = name
         self.llm = llm
+        self.name = name
         self.instruction = instruction
 
-        self.tools: list[ChatCompletionToolParam] = []
-        self.tool_functions: dict[str, Callable] = {}
+        self.tools: list[ChatCompletionToolParam] = []  # for LLM
+        
+        self.tool_functions: dict[str, Callable] = {}  # for local function call
+        self.mcp_functions: dict[str, MCPServer] = {}  # for remote function call
+        
         self.parallel_tool_calls = parallel_tool_calls
         self.use_context_variables: dict[str, str] = {}  # 使用了上下文变量的函数以及相应的形参名称
         self.use_agent_variables: dict[str, str] = {}  # 使用了Agent变量的函数以及相应的形参名称
@@ -59,6 +69,23 @@ 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 tool in tools:
+                self.tools.append({
+                    'type': 'function',
+                    'function': {
+                        'name': tool.name,
+                        'description': tool.description,
+                        'parameters': tool.inputSchema
+                    }
+                })  # 注意不能使用 add_tools 方法
+                self.mcp_functions[tool.name] = server
 
     def chat(self, message: str = None) -> ChatCompletionMessage:
         """ Agent 多轮对话
@@ -123,11 +150,13 @@ def add_tool_call_message(self, tool_content: str, tool_call_id: str) -> None:
         }
         self.messages.append(message)
 
-    def tool(self, function: Callable) -> Callable:
+    def tool(self) -> Callable:
         """ 标记一个外部工具函数
         """
-        self.add_tool_functions(function)
-        return function
+        def wrapper(function: Callable) -> Callable:
+            self.add_tool_functions(function)
+            return function
+        return wrapper
 
     def add_tools(self, *tools: 'Tool') -> 'Agent':
         """ 添加外部工具

src/course_graph/agent/mcp.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Create Date: 2025/03/28
+# Author: wangtao <wangtao.cpu@gmail.com>
+# File Name: course_graph/agent/mcp.py
+# Description: MCP Server 和 Client 实现相关
+
+from typing import Literal
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from contextlib import AsyncExitStack
+from mcp.types import Tool
+
+class MCPServer:
+    def __init__(self, type: Literal['stdio'], command: str, args: list[str], envs: dict[str, str] = None):
+        """ MCP 服务器
+
+        Args:
+            type (Literal['stdio']): 服务器类型, 目前只支持 'stdio'
+            command (str): 命令
+            args (list[str]): 参数
+            envs (dict[str, str], optional): 环境变量. Defaults to None.
+        """
+        self.params = StdioServerParameters(command=command, args=args, envs=envs)
+        self.stack = AsyncExitStack()
+        self.session = 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))
+        await self.session.initialize()
+        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()
+        
+        
+        
+        
+        

src/course_graph/agent/types.py

@@ -50,49 +50,55 @@ class Result:
     content: str = 'Function call successfully.'
     context_variables: ContextVariables | dict = ContextVariables()
     message: bool = True
-    
+
     def __repr__(self):
         return str({
             'agent': self.agent.name,
             'content': self.content,
             'context_variables': self.context_variables,
             'message': self.message
         })
-    
+
 
 @dataclass
 class TraceEvent:
     timestamp: float
     agent_name: str
 
+
 @dataclass
 class TraceEventUserMessage(TraceEvent):
     message: str
 
+
 @dataclass
 class TraceEventAgentMessage(TraceEvent):
     message: str
 
+
 @dataclass
 class TraceEventAgentSwitch(TraceEvent):
     to_agent: str
-    
+
+
 @dataclass
 class TraceEventToolCall(TraceEvent):
     function: str
     arguments: Dict[str, Any]
 
+
 @dataclass
 class TraceEventToolResult(TraceEvent):
     function: str
     result: Any
 
+
 @dataclass
 class TraceEventContextUpdate(TraceEvent):
     old_context: ContextVariables
     new_context: ContextVariables
 
-  
+
 class Trace(TypedDict):
     trace_id: str
     events: List[TraceEvent]