autogen/agents/dev_agent.py

"""
Dev Agent - 开发工程师智能体
负责代码实现和测试驱动开发
"""
from autogen import AssistantAgent
from typing import Dict, Any, Optional, List
import os
from pathlib import Path

from config.llm_config import get_agent_llm_config, DEV_PROMPT


class DevAgent:
    """开发工程师 Agent，负责编写高质量代码"""
    
    def __init__(self, llm_config: Optional[Dict] = None):
        """
        初始化 Dev Agent
        
        Args:
            llm_config: LLM 配置
        """
        self.llm_config = llm_config or get_agent_llm_config("Dev_Agent")
        
        self.agent = AssistantAgent(
            name="Dev_Agent",
            system_message=DEV_PROMPT,
            llm_config=self.llm_config,
            description="资深软件工程师，专注于汽车嵌入式 C++/Python 开发",
            human_input_mode="NEVER"
        )
        
        self.workspace_dir = Path("workspace")
        self.workspace_dir.mkdir(exist_ok=True)
        
        # 代码历史用于迭代修复
        self.code_history: List[str] = []
    
    def generate_code(self, srs_content: str, test_code: str) -> str:
        """
        根据 SRS 和测试用例生成实现代码
        
        Args:
            srs_content: SRS 文档内容
            test_code: 测试代码
        
        Returns:
            生成的源代码内容
        """
        prompt = f"""
请根据以下 SRS 和测试用例编写实现代码：

【SRS 需求】
{self._truncate(srs_content, 2500)}

【测试用例】
{self._truncate(test_code, 2000)}

请编写：
1. 完整的实现代码
2. 符合 MISRA-C/PEP8 规范
3. 包含详细的 docstring 和类型注解
4. 确保所有测试用例通过

输出 Python 代码，保存为 src_battery_health.py。
"""
        
        response = self.agent.generate_reply(
            messages=[{"role": "user", "content": prompt}]
        )
        
        code_content = response if isinstance(response, str) else str(response)
        
        # 提取代码块（如果包含 Markdown 标记）
        code_content = self._extract_code(code_content)
        
        # 保存到文件
        code_file = self.workspace_dir / "src_battery_health.py"
        with open(code_file, 'w', encoding='utf-8') as f:
            f.write(code_content)
        
        self.code_history.append(code_content)
        print(f"✅ 源代码已生成：{code_file}")
        return code_content
    
    def fix_code(self, srs_content: str, test_code: str, error_log: str) -> str:
        """
        根据测试错误修复代码
        
        Args:
            srs_content: SRS 文档
            test_code: 测试代码
            error_log: 测试失败日志
        
        Returns:
            修复后的代码
        """
        previous_code = self.code_history[-1] if self.code_history else ""
        
        prompt = f"""
代码测试失败，请根据错误日志修复代码：

【SRS 需求】
{self._truncate(srs_content, 1500)}

【测试代码】
{self._truncate(test_code, 1500)}

【之前的代码】
{self._truncate(previous_code, 2000)}

【错误日志】
{self._truncate(error_log, 1500)}

请分析错误原因并修复代码，确保：
1. 所有测试用例通过
2. 保持代码质量
3. 不引入新的问题

输出完整的修复后代码。
"""
        
        response = self.agent.generate_reply(
            messages=[{"role": "user", "content": prompt}]
        )
        
        fixed_code = response if isinstance(response, str) else str(response)
        fixed_code = self._extract_code(fixed_code)
        
        # 更新文件
        code_file = self.workspace_dir / "src_battery_health.py"
        with open(code_file, 'w', encoding='utf-8') as f:
            f.write(fixed_code)
        
        self.code_history.append(fixed_code)
        print(f"✅ 代码已修复：{code_file}")
        return fixed_code
    
    def optimize_code(self, code: str, optimization_goal: str) -> str:
        """
        优化现有代码
        
        Args:
            code: 原始代码
            optimization_goal: 优化目标（性能、可读性等）
        
        Returns:
            优化后的代码
        """
        prompt = f"""
请优化以下代码：

【原始代码】
{code}

【优化目标】
{optimization_goal}

请保持功能不变，提升代码质量。
"""
        
        response = self.agent.generate_reply(
            messages=[{"role": "user", "content": prompt}]
        )
        
        optimized_code = response if isinstance(response, str) else str(response)
        optimized_code = self._extract_code(optimized_code)
        
        return optimized_code
    
    def add_documentation(self, code: str) -> str:
        """
        为代码添加完整文档
        
        Args:
            code: 代码
        
        Returns:
            带完整文档的代码
        """
        prompt = f"""
请为以下代码添加完整的文档：

【代码】
{code}

请添加：
1. 模块级 docstring
2. 类 docstring
3. 函数 docstring（包含参数说明和返回值）
4. 类型注解
5. 关键步骤的注释

保持代码功能不变。
"""
        
        response = self.agent.generate_reply(
            messages=[{"role": "user", "content": prompt}]
        )
        
        documented_code = response if isinstance(response, str) else str(response)
        documented_code = self._extract_code(documented_code)
        
        return documented_code
    
    def _extract_code(self, text: str) -> str:
        """从文本中提取代码（移除 Markdown 标记）"""
        lines = text.split('\n')
        code_lines = []
        in_code_block = False
        
        for line in lines:
            if line.strip().startswith('```python'):
                in_code_block = True
                continue
            elif line.strip().startswith('```'):
                in_code_block = False
                continue
            
            if in_code_block or not any(line.strip().startswith(x) for x in ['```']):
                code_lines.append(line)
        
        # 如果没有找到代码块，返回原文
        if not code_lines:
            return text
        
        return '\n'.join(code_lines)
    
    def _truncate(self, text: str, max_length: int) -> str:
        """截断文本"""
        if len(text) <= max_length:
            return text
        return text[:max_length] + "... [内容已截断]"


def create_dev_agent(llm_config: Optional[Dict] = None) -> AssistantAgent:
    """
    创建 Dev Agent（AutoGen 原生格式）
    
    Args:
        llm_config: LLM 配置
    
    Returns:
        AutoGen AssistantAgent 实例
    """
    config = llm_config or get_agent_llm_config("Dev_Agent")
    
    agent = AssistantAgent(
        name="Dev_Agent",
        system_message=DEV_PROMPT,
        llm_config=config,
        description="资深软件工程师",
        human_input_mode="NEVER"
    )
    
    return agent