05 - Codex 高级功能与集成

from mcp import MCPServer, Tool, Resource

class FileSystemMCPServer(MCPServer):
    """文件系统 MCP 服务器"""
    
    def __init__(self):
        super().__init__(name="filesystem", version="1.0.0")
        
        # 注册工具
        self.register_tools()
        
        # 注册资源
        self.register_resources()
    
    def register_tools(self):
        """注册工具"""
        
        # 工具 1: 读取文件
        @self.tool(
            name="read_file",
            description="读取文件内容"
        )
        async def read_file(file_path: str) -> str:
            """
            读取文件内容
            
            Args:
                file_path: 文件路径
            
            Returns:
                文件内容
            """
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                return content
            except Exception as e:
                return f"错误：{str(e)}"
        
        # 工具 2: 写入文件
        @self.tool(
            name="write_file",
            description="写入文件内容"
        )
        async def write_file(file_path: str, content: str) -> str:
            """
            写入文件内容
            
            Args:
                file_path: 文件路径
                content: 文件内容
            
            Returns:
                执行结果
            """
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(content)
                return f"成功写入文件：{file_path}"
            except Exception as e:
                return f"错误：{str(e)}"
        
        # 工具 3: 列出目录
        @self.tool(
            name="list_directory",
            description="列出目录内容"
        )
        async def list_directory(dir_path: str) -> list:
            """
            列出目录内容
            
            Args:
                dir_path: 目录路径
            
            Returns:
                文件列表
            """
            import os
            try:
                files = os.listdir(dir_path)
                return files
            except Exception as e:
                return [f"错误：{str(e)}"]
    
    def register_resources(self):
        """注册资源"""
        
        @self.resource(
            uri="file:///{path}",
            name="File Resource",
            description="访问文件系统中的文件"
        )
        async def get_file_resource(path: str) -> Resource:
            """获取文件资源"""
            try:
                with open(path, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                return Resource(
                    uri=f"file:///{path}",
                    name=path,
                    description=f"文件：{path}",
                    mimeType="text/plain",
                    text=content
                )
            except Exception as e:
                return Resource(
                    uri=f"file:///{path}",
                    name=path,
                    description=f"错误：{str(e)}",
                    mimeType="text/plain",
                    text=""
                )

# 启动服务器
if __name__ == "__main__":
    server = FileSystemMCPServer()
    server.run(host="localhost", port=3000)

from mcp import MCPServer
import sqlite3

class DatabaseMCPServer(MCPServer):
    """数据库 MCP 服务器"""
    
    def __init__(self, db_path: str):
        super().__init__(name="database", version="1.0.0")
        self.db_path = db_path
        self.register_tools()
    
    def register_tools(self):
        """注册数据库工具"""
        
        @self.tool(
            name="execute_query",
            description="执行 SQL 查询"
        )
        async def execute_query(query: str) -> list:
            """
            执行 SQL 查询
            
            Args:
                query: SQL 查询语句
            
            Returns:
                查询结果
            """
            try:
                conn = sqlite3.connect(self.db_path)
                cursor = conn.cursor()
                cursor.execute(query)
                
                # 获取列名
                columns = [desc[0] for desc in cursor.description]
                
                # 获取数据
                rows = cursor.fetchall()
                
                # 转换为字典列表
                result = [
                    dict(zip(columns, row))
                    for row in rows
                ]
                
                conn.close()
                return result
            except Exception as e:
                return [{"error": str(e)}]
        
        @self.tool(
            name="get_schema",
            description="获取数据库结构"
        )
        async def get_schema() -> dict:
            """获取数据库结构"""
            try:
                conn = sqlite3.connect(self.db_path)
                cursor = conn.cursor()
                
                # 获取所有表
                cursor.execute(
                    "SELECT name FROM sqlite_master WHERE type='table'"
                )
                tables = cursor.fetchall()
                
                schema = {}
                
                # 获取每个表的结构
                for (table_name,) in tables:
                    cursor.execute(f"PRAGMA table_info({table_name})")
                    columns = cursor.fetchall()
                    schema[table_name] = [
                        {
                            "name": col[1],
                            "type": col[2],
                            "nullable": not col[3],
                            "primary_key": bool(col[5])
                        }
                        for col in columns
                    ]
                
                conn.close()
                return schema
            except Exception as e:
                return {"error": str(e)}

# 使用示例
server = DatabaseMCPServer(db_path="./app.db")
server.run(host="localhost", port=3001)

from openai import OpenAI
import json

client = OpenAI(api_key="your-api-key")

# 定义可用的 MCP 工具
tools = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "读取文件内容",
            "parameters": {
                "type": "object",
                "properties": {
                    "file_path": {
                        "type": "string",
                        "description": "文件路径"
                    }
                },
                "required": ["file_path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "execute_query",
            "description": "执行数据库查询",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "SQL 查询语句"
                    }
                },
                "required": ["query"]
            }
        }
    }
]

# 调用模型
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {
            "role": "system",
            "content": "你是一个编程助手，可以访问文件系统和数据库。"
        },
        {
            "role": "user",
            "content": "读取 config.py 文件，然后查询数据库中的用户表"
        }
    ],
    tools=tools,
    tool_choice="auto"
)

# 处理工具调用
if response.choices[0].message.tool_calls:
    for tool_call in response.choices[0].message.tool_calls:
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)
        
        # 调用对应的 MCP 工具
        if function_name == "read_file":
            result = call_mcp_tool("filesystem", "read_file", arguments)
        elif function_name == "execute_query":
            result = call_mcp_tool("database", "execute_query", arguments)
        
        print(f"工具调用结果：{result}")

class SecureMCPServer(MCPServer):
    """安全的 MCP 服务器"""
    
    def __init__(self):
        super().__init__()
        
        # 配置访问控制
        self.allowed_paths = ["/home/user/project"]
        self.blocked_operations = ["rm", "delete", "drop"]
    
    @self.tool(name="read_file")
    async def read_file(file_path: str) -> str:
        """安全的文件读取"""
        
        # 1. 路径验证
        if not self.is_path_allowed(file_path):
            return "错误：无权访问该路径"
        
        # 2. 路径遍历防护
        if ".." in file_path:
            return "错误：无效的文件路径"
        
        # 3. 文件大小限制
        if os.path.getsize(file_path) > 10 * 1024 * 1024:  # 10MB
            return "错误：文件过大"
        
        # 4. 读取文件
        try:
            with open(file_path, 'r') as f:
                return f.read()
        except Exception as e:
            return f"错误：{str(e)}"
    
    def is_path_allowed(self, path: str) -> bool:
        """检查路径是否允许访问"""
        abs_path = os.path.abspath(path)
        return any(
            abs_path.startswith(allowed)
            for allowed in self.allowed_paths
        )

from skills import Skill, Tool, Prompt

class PythonDevelopmentSkill(Skill):
    """Python 开发技能"""
    
    def __init__(self):
        super().__init__(
            name="Python Development",
            version="1.0.0",
            description="Python 开发的完整技能集"
        )
        
        self.register_tools()
        self.register_prompts()
        self.load_knowledge()
    
    def register_tools(self):
        """注册工具"""
        
        @self.tool(name="run_python")
        async def run_python(code: str) -> dict:
            """
            执行 Python 代码
            
            Args:
                code: Python 代码
            
            Returns:
                执行结果
            """
            import subprocess
            import tempfile
            
            try:
                # 创建临时文件
                with tempfile.NamedTemporaryFile(
                    mode='w',
                    suffix='.py',
                    delete=False
                ) as f:
                    f.write(code)
                    temp_file = f.name
                
                # 执行代码
                result = subprocess.run(
                    ['python', temp_file],
                    capture_output=True,
                    text=True,
                    timeout=5
                )
                
                return {
                    "success": result.returncode == 0,
                    "stdout": result.stdout,
                    "stderr": result.stderr
                }
            except Exception as e:
                return {
                    "success": False,
                    "error": str(e)
                }
        
        @self.tool(name="lint_python")
        async def lint_python(code: str) -> list:
            """
            检查 Python 代码风格
            
            Args:
                code: Python 代码
            
            Returns:
                问题列表
            """
            import pylint.lint
            import tempfile
            
            try:
                # 创建临时文件
                with tempfile.NamedTemporaryFile(
                    mode='w',
                    suffix='.py',
                    delete=False
                ) as f:
                    f.write(code)
                    temp_file = f.name
                
                # 运行 pylint
                results = pylint.lint.Run(
                    [temp_file],
                    do_exit=False
                )
                
                # 解析结果
                issues = []
                for msg in results.linter.reporter.messages:
                    issues.append({
                        "line": msg.line,
                        "column": msg.column,
                        "message": msg.msg,
                        "type": msg.category
                    })
                
                return issues
            except Exception as e:
                return [{"error": str(e)}]
        
        @self.tool(name="generate_tests")
        async def generate_tests(code: str) -> str:
            """
            生成测试用例
            
            Args:
                code: 要测试的代码
            
            Returns:
                测试代码
            """
            # 使用 AI 生成测试
            prompt = f"""
            为以下 Python 代码生成 pytest 测试用例：
            
            ```python
            {code}
            ```
            
            要求：
            1. 测试正常情况
            2. 测试边界条件
            3. 测试异常情况
            4. 使用清晰的测试名称
            """
            
            # 调用 AI 模型生成测试
            test_code = await self.generate_with_ai(prompt)
            return test_code
    
    def register_prompts(self):
        """注册提示词模板"""
        
        self.add_prompt(
            name="code_review",
            template="""
            请审查以下 Python 代码：
            
            ```python
            {code}
            ```
            
            请关注：
            1. 代码风格（PEP 8）
            2. 潜在 Bug
            3. 性能问题
            4. 安全隐患
            5. 可维护性
            
            给出具体的改进建议。
            """,
            variables=["code"]
        )
        
        self.add_prompt(
            name="optimize_performance",
            template="""
            优化以下 Python 代码的性能：
            
            ```python
            {code}
            ```
            
            当前性能指标：
            - 执行时间：{execution_time}ms
            - 内存使用：{memory_usage}MB
            
            请提供：
            1. 性能瓶颈分析
            2. 优化方案
            3. 优化后的代码
            4. 预期的性能提升
            """,
            variables=["code", "execution_time", "memory_usage"]
        )
    
    def load_knowledge(self):
        """加载知识库"""
        
        # 加载最佳实践文档
        self.add_knowledge(
            name="Python Best Practices",
            content="""
            # Python 最佳实践
            
            ## 1. 代码风格
            - 遵循 PEP 8
            - 使用有意义的变量名
            - 保持函数简短（< 50 行）
            
            ## 2. 性能优化
            - 使用列表推导式而非循环
            - 避免在循环中进行字符串拼接
            - 使用生成器处理大数据
            
            ## 3. 错误处理
            - 使用具体的异常类型
            - 不要捕获所有异常
            - 始终清理资源（使用 with 语句）
            
            ## 4. 测试
            - 编写单元测试
            - 使用 pytest 框架
            - 追求高测试覆盖率
            """
        )

# 使用 Skill
skill = PythonDevelopmentSkill()

# 执行代码审查
code = """
def add(a, b):
    return a+b
"""

review_result = await skill.execute_prompt(
    "code_review",
    code=code
)

print(review_result)

class CompositeSkill(Skill):
    """组合多个 Skills"""
    
    def __init__(self):
        super().__init__(name="Full Stack Development")
        
        # 加载子 Skills
        self.python_skill = PythonDevelopmentSkill()
        self.frontend_skill = FrontendDevelopmentSkill()
        self.devops_skill = DevOpsSkill()
    
    async def build_web_app(self, requirements: str):
        """构建完整的 Web 应用"""
        
        # 1. 使用 Python Skill 创建后端
        backend_code = await self.python_skill.generate_api(requirements)
        
        # 2. 使用 Frontend Skill 创建前端
        frontend_code = await self.frontend_skill.generate_ui(requirements)
        
        # 3. 使用 DevOps Skill 配置部署
        deployment_config = await self.devops_skill.generate_deployment(
            backend=backend_code,
            frontend=frontend_code
        )
        
        return {
            "backend": backend_code,
            "frontend": frontend_code,
            "deployment": deployment_config
        }

# 使用
composite_skill = CompositeSkill()
result = await composite_skill.build_web_app(
    "创建一个待办事项管理应用"
)

from typing import List, Dict
from openai import OpenAI

class CodeAgent:
    """代码生成 Agent"""
    
    def __init__(self):
        self.client = OpenAI(api_key="your-api-key")
        self.memory = []  # 对话历史
        self.tools = self.register_tools()
    
    def register_tools(self) -> List[Dict]:
        """注册可用工具"""
        return [
            {
                "type": "function",
                "function": {
                    "name": "read_file",
                    "description": "读取文件内容",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "path": {"type": "string"}
                        }
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "write_file",
                    "description": "写入文件",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "path": {"type": "string"},
                            "content": {"type": "string"}
                        }
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "run_tests",
                    "description": "运行测试",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "test_path": {"type": "string"}
                        }
                    }
                }
            }
        ]
    
    async def execute_task(self, task: str) -> str:
        """执行任务"""
        
        # 1. 规划阶段
        plan = await self.plan(task)
        print(f"📋 计划：{plan}")
        
        # 2. 执行阶段
        results = []
        for step in plan['steps']:
            result = await self.execute_step(step)
            results.append(result)
            print(f"✅ 完成步骤：{step['description']}")
        
        # 3. 反思阶段
        reflection = await self.reflect(task, results)
        print(f"💭 反思：{reflection}")
        
        return reflection['final_result']
    
    async def plan(self, task: str) -> Dict:
        """制定执行计划"""
        
        prompt = f"""
        任务：{task}
        
        请制定详细的执行计划，分解为具体步骤。
        
        返回 JSON 格式：
        {{
            "steps": [
                {{
                    "id": 1,
                    "description": "步骤描述",
                    "tool": "使用的工具",
                    "expected_outcome": "预期结果"
                }}
            ]
        }}
        """
        
        response = await self.call_llm(prompt)
        return json.loads(response)
    
    async def execute_step(self, step: Dict) -> Dict:
        """执行单个步骤"""
        
        tool_name = step['tool']
        
        # 调用对应的工具
        if tool_name == "read_file":
            result = self.read_file(step['params']['path'])
        elif tool_name == "write_file":
            result = self.write_file(
                step['params']['path'],
                step['params']['content']
            )
        elif tool_name == "run_tests":
            result = self.run_tests(step['params']['test_path'])
        
        return {
            "step_id": step['id'],
            "result": result,
            "success": result.get('success', False)
        }
    
    async def reflect(self, task: str, results: List[Dict]) -> Dict:
        """反思执行结果"""
        
        prompt = f"""
        原始任务：{task}
        
        执行结果：
        {json.dumps(results, indent=2)}
        
        请评估：
        1. 任务是否成功完成
        2. 有哪些改进空间
        3. 需要重新执行哪些步骤
        
        返回 JSON 格式：
        {{
            "task_completed": true/false,
            "improvements": ["改进建议"],
            "retry_steps": [步骤ID],
            "final_result": "最终结果描述"
        }}
        """
        
        response = await self.call_llm(prompt)
        return json.loads(response)
    
    async def call_llm(self, prompt: str) -> str:
        """调用 LLM"""
        
        self.memory.append({"role": "user", "content": prompt})
        
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=self.memory,
            tools=self.tools
        )
        
        result = response.choices[0].message.content
        self.memory.append({"role": "assistant", "content": result})
        
        return result

# 使用示例
agent = CodeAgent()

result = await agent.execute_task(
    "创建一个 REST API，包含用户增删改查功能，并编写测试"
)

print(f"最终结果：{result}")