scripts/02_generate_data.py

"""
修复版训练数据生成器
核心改进:
1. 直接基于代码内容生成准确的问答对
2. 不依赖LLM生成（避免循环依赖）
3. 使用模板化方法确保数据质量
4. 优化项目概览问题，使其更具项目特色
"""
import json
import yaml
import random
from pathlib import Path
from typing import List, Dict, Any
from dataclasses import dataclass, field # <--- 修复: dataclass 位于 dataclasses 模块
import re
from collections import defaultdict


@dataclass
class TrainingSample:
    """训练样本"""
    conversations: List[Dict[str, str]]
    metadata: Dict[str, Any]


class FixedDataGenerator:
    """修复版数据生成器 - 基于规则和模板"""
    
    def __init__(self, config_path: str = "../config/default_config.yaml", 
                 analysis_path: str = "../data/repository_analysis.json"):
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        try:
            with open(analysis_path, 'r', encoding='utf-8') as f:
                self.analysis_data = json.load(f)
        except FileNotFoundError:
            print(f"❌ 警告: 找不到分析文件 {analysis_path}。请先运行分析器。")
            self.analysis_data = {'code_elements': [], 'project_context': {}}

        self.code_elements = self.analysis_data.get('code_elements', [])
        self.project_context = self.analysis_data.get('project_context', {})
        self.project_name = self.project_context.get('project_name', 'Laddr')
        
        self.training_samples = []
    
    def generate_training_data(self):
        """生成训练数据"""
        print(f"Generating training data for {self.project_name}...")
        
        # 1. 代码解释任务（基于docstring + 代码结构）
        print("Generating code explanation samples...")
        self._generate_code_explanation_samples()
        
        # 2. API使用示例（基于函数签名 + docstring）
        print("Generating API usage samples...")
        self._generate_api_usage_samples()
        
        # 3. 项目概览问答（基于统计和结构信息）
        print("Generating project overview samples...")
        self._generate_project_overview_samples()
        
        # 4. 代码定位任务（"在哪个文件中..."）
        print("Generating code location samples...")
        self._generate_code_location_samples()
        
        print(f"Total samples generated: {len(self.training_samples)}")
    
    def _generate_code_explanation_samples(self):
        """生成代码解释样本 - 基于真实代码和docstring"""
        # 选择有docstring的元素
        candidates = [e for e in self.code_elements 
                     if e.get('docstring') and len(e.get('code', '')) > 50]
        
        for element in candidates[:300]:  # 增加数量限制
            name = element['name']
            docstring = element['docstring']
            filepath = element['filepath']
            element_type = element['type']
            code = element.get('code', '')
            
            # 提取函数签名
            signature = self._extract_signature(code, element_type)
            
            # 问题模板
            questions = [
                f"请解释 {self.project_name} 中 `{name}` 的作用。",
                f"{self.project_name} 的 `{name}` 是做什么的？",
                f"在 {self.project_name} 项目中，`{name}` 有什么功能？",
            ]
            question = random.choice(questions)
            
            # 构建高质量答案（基于真实信息）
            answer_parts = []
            
            # 1. 基本信息
            answer_parts.append(f"`{name}` 是 {self.project_name} 项目中的一个 {self._type_to_cn(element_type)}，位于 `{filepath}`。")
            
            # 2. 功能描述（来自docstring）
            if docstring:
                # 清理docstring
                clean_doc = self._clean_docstring(docstring)
                answer_parts.append(f"\n**功能描述**：\n{clean_doc}")
            
            # 3. 函数签名（如果有）
            if signature:
                answer_parts.append(f"\n**函数签名**：\n```python\n{signature}\n```")
            
            # 4. 参数说明（如果有）
            params = element.get('parameters', [])
            if params and len(params) > 0:
                param_desc = "\n**参数**：\n"
                for param in params[:5]:  # 最多5个参数
                    param_name = param.get('name', 'unknown')
                    param_type = param.get('type', 'Any')
                    # 尝试从 docstring 中提取参数描述，如果没有则使用类型
                    param_desc_from_doc = self._extract_param_desc(docstring, param_name)
                    if param_desc_from_doc:
                        param_info = f"- `{param_name}` ({param_type}): {param_desc_from_doc}\n"
                    else:
                        param_info = f"- `{param_name}` ({param_type})\n"

                    param_desc += param_info
                answer_parts.append(param_desc)
            
            # 5. 返回值（如果有）
            return_type = element.get('return_type')
            if return_type:
                answer_parts.append(f"\n**返回值**：`{return_type}`")
            
            answer = ''.join(answer_parts)
            
            self.training_samples.append(TrainingSample(
                conversations=[
                    {"role": "user", "content": question},
                    {"role": "assistant", "content": answer}
                ],
                metadata={
                    "task_type": "code_explanation",
                    "element_name": name,
                    "filepath": filepath
                }
            ))
    
    def _generate_api_usage_samples(self):
        """生成API使用示例 - 基于函数签名"""
        # 选择公共函数/方法
        candidates = [e for e in self.code_elements 
                     if e['type'] in ['function', 'method']
                     and not e['name'].startswith('_')  # 排除私有方法
                     and e.get('parameters')]
        
        for element in candidates[:150]: # 增加数量限制
            name = element['name']
            params = element.get('parameters', [])
            filepath = element['filepath']
            docstring = element.get('docstring', '')
            
            question = f"如何在 {self.project_name} 中使用 `{name}` 函数？"
            
            # 构建使用示例
            answer_parts = []
            answer_parts.append(f"`{name}` 位于 `{filepath}`，使用方法如下：")
            
            # 生成示例代码
            param_names = [p['name'] for p in params if p['name'] != 'self']
            if param_names:
                example_code = f"{name}("
                param_examples = []
                for p in param_names[:5]:  # 最多5个参数
                    param_examples.append(f"{p}=...")
                example_code += ", ".join(param_examples)
                example_code += ")"
                
                answer_parts.append(f"\n```python\n{example_code}\n```")
            
            # 参数说明
            if params:
                answer_parts.append("\n**参数说明**：")
                for param in params[:5]:
                    if param['name'] != 'self':
                        param_type = param.get('type', 'Any')
                        
                        param_desc_from_doc = self._extract_param_desc(docstring, param['name'])
                        
                        answer_parts.append(f"\n- `{param['name']}`: {param_type}")
                        if param_desc_from_doc:
                             answer_parts[-1] += f" - {param_desc_from_doc}" # 追加描述
            
            # 添加docstring提示
            if docstring:
                clean_doc = self._clean_docstring(docstring)[:200]
                if clean_doc:
                    answer_parts.append(f"\n\n**功能简述**：{clean_doc}...")
            
            answer = ''.join(answer_parts)
            
            self.training_samples.append(TrainingSample(
                conversations=[
                    {"role": "user", "content": question},
                    {"role": "assistant", "content": answer}
                ],
                metadata={
                    "task_type": "api_usage",
                    "element_name": name
                }
            ))
    
    def _generate_project_overview_samples(self):
        """生成项目概览问答 - 基于统计信息"""
        stats = self.analysis_data.get('statistics', {})
        description = self.project_context.get('description', '')
        techs = self.project_context.get('main_technologies', [])
        file_type_counts = self.analysis_data.get('statistics', {}).get('file_type_counts', {})
        
        # --- 问题1: 项目主要功能 (更具项目特色) ---
        q1_list = [
            f"请用一句话描述 {self.project_name} 项目的主要功能。",
            f"{self.project_name} 是一个什么样的项目？",
            f"简单介绍一下 {self.project_name} 项目。"
        ]
        q1 = random.choice(q1_list)
        
        a1_parts = [
            f"{self.project_name} 是一个 Python 项目。"
        ]
        
        if description:
            # 修复：确保项目描述清晰
            a1_parts.append(f"\n**核心目标**：\n{description}")
        else:
             a1_parts.append("\n**核心目标**：此项目旨在提供一个可扩展的多代理系统框架（Agent Framework），支持任务规划、工具调用、消息队列和数据库集成等功能。")
        
        # 添加技术栈
        if techs:
            a1_parts.append(f"\n\n**主要技术栈**：{', '.join(techs[:5])}等。")
        
        a1_parts.append(f"\n\n项目包含 {stats.get('total_elements', 0)} 个代码元素，主要由 {stats.get('classes', 0)} 个类和 {stats.get('functions', 0) + stats.get('methods', 0)} 个函数/方法构成。")
        
        a1 = ''.join(a1_parts)
        
        self.training_samples.append(TrainingSample(
            conversations=[
                {"role": "user", "content": q1},
                {"role": "assistant", "content": a1}
            ],
            metadata={"task_type": "project_overview"}
        ))
        
        # --- 问题2: 项目结构 ---
        q2_list = [
            f"{self.project_name} 的项目结构是怎样的？",
            f"请列举 {self.project_name} 的核心模块。",
        ]
        q2 = random.choice(q2_list)
        
        a2_parts = [f"{self.project_name} 项目包含以下主要部分：\n"]
        
        # 获取主要模块
        modules = self.project_context.get('key_modules', [])
        if modules:
            a2_parts.append("\n**核心模块**：\n")
            for mod in modules[:10]:
                a2_parts.append(f"- `{mod}`\n")
        else:
            a2_parts.append("\n**核心模块**：\n- `core` (核心逻辑，如Agent Runtime, Tooling, Config)\n- `cli` (命令行接口)\n- `llms` (LLM后端实现)\n")

        # 优化文件类型展示
        if file_type_counts:
             file_stats = ', '.join(f'{k.lstrip(".").upper()}: {v}' for k, v in file_type_counts.items() if k not in ['.other'])
             a2_parts.append(f"\n**主要文件类型统计**：{file_stats}")
        
        a2 = ''.join(a2_parts)
        
        self.training_samples.append(TrainingSample(
            conversations=[
                {"role": "user", "content": q2},
                {"role": "assistant", "content": a2}
            ],
            metadata={"task_type": "project_structure"}
        ))
        
        # --- 问题3: 核心类/函数 ---
        top_elements = sorted(self.code_elements, 
                             key=lambda x: x.get('complexity', 0), 
                             reverse=True)[:10]
        
        q3 = f"{self.project_name} 中有哪些核心类和函数？"
        a3_parts = [f"{self.project_name} 的核心组件包括（基于复杂度和重要性）：\n"]
        
        for elem in top_elements:
            name = elem['name']
            filepath = elem['filepath']
            elem_type = self._type_to_cn(elem['type'])
            
            doc = elem.get('docstring', '')
            short_doc = self._clean_docstring(doc).split('\n')[0][:80].strip()
            
            line = f"\n- `{name}` ({elem_type})：位于 `{filepath}`"
            if short_doc:
                line += f" - {short_doc}..."
            a3_parts.append(line)
        
        if len(top_elements) > 0:
             a3 = ''.join(a3_parts)
             self.training_samples.append(TrainingSample(
                conversations=[
                    {"role": "user", "content": q3},
                    {"role": "assistant", "content": a3}
                ],
                metadata={"task_type": "core_components"}
            ))
    
    def _generate_code_location_samples(self):
        """生成代码定位任务"""
        # 选择不同文件中的元素
        file_elements = defaultdict(list)
        for elem in self.code_elements:
            # 排除非核心的__init__
            if elem['name'] == '__init__' and 'module' not in elem['type']:
                continue
            file_elements[elem['filepath']].append(elem)
        
        for filepath, elements in list(file_elements.items())[:50]:
            # 随机选择1-3个元素
            selected = random.sample(elements, min(3, len(elements)))
            
            for elem in selected:
                name = elem['name']
                elem_type = self._type_to_cn(elem['type'])
                
                question = f"在 {self.project_name} 中，`{name}` {elem_type}在哪个文件里？"
                
                # 答案优化：更简洁，减少冗余信息，模型只需学习路径
                answer = f"`{name}` 位于 `{filepath}`。"
                
                self.training_samples.append(TrainingSample(
                    conversations=[
                        {"role": "user", "content": question},
                        {"role": "assistant", "content": answer}
                    ],
                    metadata={
                        "task_type": "code_location",
                        "element_name": name,
                        "filepath": filepath
                    }
                ))
    
    def _extract_signature(self, code: str, element_type: str) -> str:
        """提取函数/类签名"""
        if not code:
            return ""
        
        lines = code.strip().split('\n')
        signature_lines = []
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            signature_lines.append(line)
            
            # 提取函数/方法定义行
            if element_type in ['function', 'method'] and (line.startswith('def ') or line.startswith('async def ')):
                # 兼容多行函数签名
                if not line.endswith(':'):
                    continue
                return '\n'.join(signature_lines)
            
            # 提取类定义行
            if element_type == 'class' and line.startswith('class '):
                 if not line.endswith(':'):
                    continue
                 return '\n'.join(signature_lines)

            # 避免包含函数/方法体
            if line.endswith((':')) and not line.startswith(('def ', 'class ')):
                break

        # 仅返回前几行，确保只包含定义
        return '\n'.join(signature_lines[:5])
    
    def _clean_docstring(self, docstring: str) -> str:
        """清理docstring"""
        if not docstring:
            return ""
        
        # 移除多余空白
        lines = docstring.strip().split('\n')
        cleaned = []
        for line in lines:
            line = line.strip()
            if line:
                cleaned.append(line)
        
        return ' '.join(cleaned)

    def _extract_param_desc(self, docstring: str, param_name: str) -> str:
        """从 docstring 中尝试提取参数描述"""
        if not docstring:
            return ""
        # 匹配各种格式的参数描述，例如 Args: key: The cache key.
        match = re.search(rf"(?:Args|Parameters|Params):\s*(?:[\n\r]\s*-)?\s*`?{re.escape(param_name)}`?\s*[:\-]\s*(.*)", docstring, re.IGNORECASE)
        if match:
             desc = match.group(1).split('\n')[0].strip()
             return desc if desc else "无描述"
        return ""
    
    def _type_to_cn(self, element_type: str) -> str:
        """元素类型转中文"""
        mapping = {
            'function': '函数',
            'method': '方法',
            'class': '类',
            'variable': '变量',
            'module': '模块'
        }
        return mapping.get(element_type, element_type)
    
    def save_training_data(self):
        """保存训练数据"""
        output_dir = Path(self.config['dataset']['output_dir'])
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 打乱
        random.shuffle(self.training_samples)
        
        # 分割
        total = len(self.training_samples)
        train_size = int(total * 0.8)
        val_size = int(total * 0.1)
        
        if total < 10: # 如果样本太少，平均分配
             train_size = max(1, total // 2)
             val_size = max(1, (total - train_size) // 2)
        
        # 再次检查，确保分割不会导致索引错误
        if train_size + val_size > total:
             val_size = total - train_size

        train_data = self.training_samples[:train_size]
        val_data = self.training_samples[train_size:train_size + val_size]
        test_data = self.training_samples[train_size + val_size:]
        
        # 保存为JSONL
        self._save_jsonl(train_data, output_dir / "train.jsonl")
        self._save_jsonl(val_data, output_dir / "val.jsonl")
        self._save_jsonl(test_data, output_dir / "test.jsonl")
        
        # 元数据
        metadata = {
            'total_samples': total,
            'train_samples': len(train_data),
            'val_samples': len(val_data),
            'test_samples': len(test_data),
            'project_name': self.project_name,
            'task_distribution': self._get_task_distribution()
        }
        
        with open(output_dir / "metadata.json", 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)
        
        print(f"\n✓ Training data saved:")
        print(f"  Train: {len(train_data)}")
        print(f"  Val: {len(val_data)}")
        print(f"  Test: {len(test_data)}")
        print(f"  Total: {total}")
        
        # 显示样本示例
        print(f"\n📝 Sample training example:")
        if train_data:
            sample = random.choice(train_data)
            print(f"Q: {sample.conversations[0]['content'][:100]}...")
            print(f"A: {sample.conversations[1]['content'][:150]}...")
    
    def _save_jsonl(self, data: List[TrainingSample], filepath: Path):
        """保存为JSONL格式"""
        with open(filepath, 'w', encoding='utf-8') as f:
            for sample in data:
                # 仅保存对话，不保存 metadata
                json.dump({'conversations': sample.conversations}, f, ensure_ascii=False)
                f.write('\n')
    
    def _get_task_distribution(self) -> Dict[str, int]:
        """统计任务分布"""
        dist = {}
        for sample in self.training_samples:
            task_type = sample.metadata.get('task_type', 'unknown')
            dist[task_type] = dist.get(task_type, 0) + 1
        return dist


def main():
    print("="*60)
    print("Fixed Training Data Generator (Project-Specific Answers Enhanced)")
    print("="*60)
    
    generator = FixedDataGenerator()
    generator.generate_training_data()
    generator.save_training_data()
    
    print("\n" + "="*60)
    print("✓ Data generation completed!")
    print("="*60)


if __name__ == "__main__":
    main()