554 lines
18 KiB
Python
554 lines
18 KiB
Python
"""
|
||
Neural Sync - Code Draft Generation Service
|
||
|
||
AI 代码草案生成服务 - 激进策略
|
||
AI 尝试生成所有代码（包括复杂逻辑），但提供详细的 review points
|
||
"""
|
||
|
||
import asyncio
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import logging
|
||
from typing import List, Optional, Dict, Any
|
||
from pydantic import BaseModel, Field
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from ..utils.timezone import now_cn_aware
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class CodeDraftRequest(BaseModel):
|
||
"""代码草案生成请求"""
|
||
signal_id: str
|
||
signal_content: str
|
||
strategic_analysis: Dict[str, Any]
|
||
target_files: List[str]
|
||
project_root: str = "."
|
||
|
||
|
||
class FileInfo(BaseModel):
|
||
"""文件信息"""
|
||
path: str
|
||
language: str
|
||
exists: bool
|
||
content: Optional[str] = None
|
||
lines_count: int = 0
|
||
|
||
|
||
class GeneratedFile(BaseModel):
|
||
"""生成的文件"""
|
||
path: str
|
||
content: str
|
||
language: str
|
||
diff: str
|
||
is_new: bool
|
||
review_points: List[str] = []
|
||
|
||
|
||
class CodeDraftResponse(BaseModel):
|
||
"""代码草案生成响应"""
|
||
success: bool
|
||
files: List[GeneratedFile]
|
||
summary: str
|
||
confidence: float = Field(default=0.0, ge=0.0, le=1.0)
|
||
warnings: List[str] = Field(default_factory=list)
|
||
review_points: List[str] = Field(default_factory=list)
|
||
estimated_complexity: str = "low"
|
||
generation_time: float = 0.0
|
||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||
|
||
|
||
class CodeDraftService:
|
||
"""
|
||
AI 代码草案生成服务 - 激进策略
|
||
|
||
特性:
|
||
- 支持全栈代码生成
|
||
- 读取现有代码上下文
|
||
- 生成 unified diff
|
||
- 提取人工 review 关键点
|
||
- 深度代码验证
|
||
"""
|
||
|
||
def __init__(self, llm_service):
|
||
self.llm = llm_service
|
||
self._language_map = {
|
||
".py": "python",
|
||
".ts": "typescript",
|
||
".tsx": "typescript",
|
||
".js": "javascript",
|
||
".jsx": "javascript",
|
||
".vue": "vue",
|
||
".sql": "sql",
|
||
".sh": "bash",
|
||
".yaml": "yaml",
|
||
".yml": "yaml",
|
||
".json": "json",
|
||
}
|
||
|
||
async def generate_draft(self, request: CodeDraftRequest) -> CodeDraftResponse:
|
||
"""
|
||
生成代码草案（激进策略）
|
||
"""
|
||
start_time = now_cn_aware().timestamp()
|
||
|
||
try:
|
||
# 1. 读取目标文件
|
||
file_contexts = await self._read_target_files(request)
|
||
|
||
# 2. 分析代码库结构
|
||
repo_context = await self._analyze_repo_structure(request)
|
||
|
||
# 3. 构建激进策略的生成 prompt
|
||
prompt = self._build_aggressive_prompt(request, file_contexts, repo_context)
|
||
|
||
# 4. 调用 LLM 生成代码
|
||
llm_response = await self._call_llm_for_code(prompt)
|
||
|
||
# 5. 解析和验证生成的代码
|
||
generated_files = await self._parse_and_validate(
|
||
llm_response, file_contexts
|
||
)
|
||
|
||
# 6. 生成 unified diff
|
||
for file in generated_files:
|
||
file.diff = self._generate_diff(file_contexts, file)
|
||
|
||
# 7. 提取 review points
|
||
review_points = await self._extract_review_points(
|
||
generated_files, file_contexts, request
|
||
)
|
||
|
||
# 8. 计算置信度和复杂度
|
||
confidence = self._calculate_confidence(generated_files, file_contexts)
|
||
complexity = self._estimate_complexity(generated_files)
|
||
|
||
# 9. 生成摘要
|
||
summary = self._generate_summary(generated_files, request)
|
||
|
||
# 10. 收集警告
|
||
warnings = self._collect_warnings(generated_files, file_contexts)
|
||
|
||
generation_time = now_cn_aware().timestamp() - start_time
|
||
|
||
return CodeDraftResponse(
|
||
success=True,
|
||
files=generated_files,
|
||
summary=summary,
|
||
confidence=confidence,
|
||
warnings=warnings,
|
||
review_points=review_points,
|
||
estimated_complexity=complexity,
|
||
generation_time=generation_time,
|
||
metadata={
|
||
"signal_id": request.signal_id,
|
||
"files_count": len(generated_files),
|
||
"new_files": sum(1 for f in generated_files if f.is_new),
|
||
"modified_files": sum(1 for f in generated_files if not f.is_new),
|
||
},
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Code generation failed: {e}", exc_info=True)
|
||
return CodeDraftResponse(
|
||
success=False,
|
||
files=[],
|
||
summary=f"代码生成失败: {str(e)}",
|
||
confidence=0.0,
|
||
warnings=[f"生成过程出错: {str(e)}"],
|
||
review_points=[],
|
||
generation_time=now_cn_aware().timestamp() - start_time,
|
||
)
|
||
|
||
async def _read_target_files(
|
||
self, request: CodeDraftRequest
|
||
) -> Dict[str, FileInfo]:
|
||
"""读取目标文件内容"""
|
||
file_contexts = {}
|
||
|
||
for file_path in request.target_files:
|
||
full_path = os.path.join(request.project_root, file_path)
|
||
|
||
info = FileInfo(
|
||
path=file_path,
|
||
language=self._detect_language(file_path),
|
||
exists=os.path.exists(full_path),
|
||
)
|
||
|
||
if info.exists:
|
||
try:
|
||
with open(full_path, "r", encoding="utf-8") as f:
|
||
info.content = f.read()
|
||
info.lines_count = len(info.content.split("\n"))
|
||
except Exception as e:
|
||
logger.warning(f"Failed to read {file_path}: {e}")
|
||
info.content = None
|
||
|
||
file_contexts[file_path] = info
|
||
|
||
return file_contexts
|
||
|
||
def _detect_language(self, file_path: str) -> str:
|
||
"""检测文件语言"""
|
||
for ext, lang in self._language_map.items():
|
||
if file_path.endswith(ext):
|
||
return lang
|
||
return "text"
|
||
|
||
async def _analyze_repo_structure(self, request: CodeDraftRequest) -> Dict[str, Any]:
|
||
"""分析代码库结构"""
|
||
repo_context = {
|
||
"existing_modules": [],
|
||
"imports_available": [],
|
||
"patterns_used": [],
|
||
}
|
||
|
||
try:
|
||
# 扫描 Python 文件的导入
|
||
py_files = list(
|
||
Path(request.project_root).rglob("*.py")
|
||
)
|
||
for py_file in py_files[:20]: # 限制扫描数量
|
||
try:
|
||
content = py_file.read_text(encoding="utf-8", errors="ignore")
|
||
# 提取导入
|
||
imports = re.findall(r"^import (\w+)|^from (\w+) import", content, re.MULTILINE)
|
||
for imp in imports:
|
||
module = imp[0] or imp[1]
|
||
if module and module not in repo_context["imports_available"]:
|
||
repo_context["imports_available"].append(module)
|
||
except:
|
||
pass
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Repo structure analysis failed: {e}")
|
||
|
||
return repo_context
|
||
|
||
def _build_aggressive_prompt(
|
||
self,
|
||
request: CodeDraftRequest,
|
||
file_contexts: Dict[str, FileInfo],
|
||
repo_context: Dict[str, Any],
|
||
) -> str:
|
||
"""构建激进的代码生成 prompt"""
|
||
prompt = f"""You are a senior full-stack engineer. Generate COMPLETE, PRODUCTION-READY code based on the following requirements.
|
||
|
||
## Signal (Requirement)
|
||
**Content**: {request.signal_content}
|
||
|
||
## Strategic Analysis
|
||
```json
|
||
{self._format_json(request.strategic_analysis)}
|
||
```
|
||
|
||
## Available Imports/Modules
|
||
{', '.join(repo_context['imports_available'][:20])}
|
||
|
||
## Existing Code Context
|
||
"""
|
||
|
||
for file_path, info in file_contexts.items():
|
||
prompt += f"\n### {file_path} ({info.language})\n"
|
||
if info.exists and info.content:
|
||
# 只显示关键部分（导入、类定义、函数签名）
|
||
lines = info.content.split("\n")
|
||
key_lines = []
|
||
for line in lines:
|
||
stripped = line.strip()
|
||
# 保留导入、类/函数定义、装饰器
|
||
if any(
|
||
stripped.startswith(p)
|
||
for p in ["import ", "from ", "class ", "def ", "@", "async def "]
|
||
):
|
||
key_lines.append(line)
|
||
elif len(key_lines) < 30: # 保留前30行作为上下文
|
||
key_lines.append(line)
|
||
|
||
prompt += "```" + info.language + "\n"
|
||
prompt += "\n".join(key_lines)
|
||
if len(lines) > 30:
|
||
prompt += f"\n... ({len(lines) - 30} more lines)\n"
|
||
prompt += "```\n"
|
||
else:
|
||
prompt += "*(New file to be created)*\n"
|
||
|
||
prompt += """
|
||
## Code Generation Requirements
|
||
1. **Complete Implementation**: Generate FULL, WORKING code including complex business logic
|
||
2. **Error Handling**: Proper exception handling and edge cases
|
||
3. **Performance**: Consider performance implications
|
||
4. **Type Safety**: Use proper type hints
|
||
5. **Comments**: Add concise comments for complex logic
|
||
6. **Testing**: Consider testability
|
||
|
||
## Output Format
|
||
Return a JSON response with this EXACT structure:
|
||
```json
|
||
{
|
||
"files": [
|
||
{
|
||
"path": "relative/path/to/file.ext",
|
||
"content": "full file content here",
|
||
"is_new": true/false,
|
||
"review_points": ["list of key points needing human review"]
|
||
}
|
||
],
|
||
"summary": "Brief description of changes",
|
||
"warnings": ["any potential issues or risks"]
|
||
}
|
||
```
|
||
|
||
IMPORTANT:
|
||
- Generate complete, runnable code
|
||
- Handle edge cases and errors
|
||
- Follow existing code patterns in the project
|
||
- Highlight complex logic in review_points
|
||
"""
|
||
|
||
return prompt
|
||
|
||
async def _call_llm_for_code(self, prompt: str) -> Dict[str, Any]:
|
||
"""调用 LLM 生成代码"""
|
||
try:
|
||
# 使用 LLM 服务的聊天接口
|
||
response = await self.llm._call_llm_chat(
|
||
messages=[
|
||
{"role": "system", "content": "You are an expert software engineer."},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.3, # 降低温度以获得更稳定的代码
|
||
)
|
||
|
||
# 解析 JSON 响应
|
||
json_match = re.search(r'```json\s*([\s\S]*?)\s*```', response)
|
||
if json_match:
|
||
import json
|
||
return json.loads(json_match.group(1))
|
||
|
||
# 尝试直接解析 JSON
|
||
import json
|
||
return json.loads(response)
|
||
|
||
except Exception as e:
|
||
logger.error(f"LLM call or parsing failed: {e}")
|
||
# 返回错误响应
|
||
return {
|
||
"files": [],
|
||
"summary": f"LLM 生成失败: {str(e)}",
|
||
"warnings": [str(e)]
|
||
}
|
||
|
||
async def _parse_and_validate(
|
||
self, llm_response: Dict[str, Any], file_contexts: Dict[str, FileInfo]
|
||
) -> List[GeneratedFile]:
|
||
"""解析和验证生成的代码"""
|
||
generated_files = []
|
||
|
||
for file_data in llm_response.get("files", []):
|
||
file_path = file_data.get("path", "")
|
||
content = file_data.get("content", "")
|
||
review_points = file_data.get("review_points", [])
|
||
|
||
file_info = file_contexts.get(file_path)
|
||
if not file_info:
|
||
file_info = FileInfo(
|
||
path=file_path,
|
||
language=self._detect_language(file_path),
|
||
exists=False,
|
||
)
|
||
|
||
# 基础验证
|
||
validation_warnings = self._validate_code_syntax(content, file_info.language)
|
||
if validation_warnings:
|
||
review_points.extend(validation_warnings)
|
||
|
||
generated_files.append(
|
||
GeneratedFile(
|
||
path=file_path,
|
||
content=content,
|
||
language=file_info.language,
|
||
diff="", # 稍后生成
|
||
is_new=not file_info.exists,
|
||
review_points=review_points,
|
||
)
|
||
)
|
||
|
||
return generated_files
|
||
|
||
def _validate_code_syntax(self, content: str, language: str) -> List[str]:
|
||
"""验证代码语法"""
|
||
warnings = []
|
||
|
||
try:
|
||
if language == "python":
|
||
import ast
|
||
ast.parse(content)
|
||
elif language in ["typescript", "javascript"]:
|
||
# 简单的括号匹配检查
|
||
if content.count("(") != content.count(")"):
|
||
warnings.append("括号不匹配")
|
||
if content.count("{") != content.count("}"):
|
||
warnings.append("花括号不匹配")
|
||
except SyntaxError as e:
|
||
warnings.append(f"语法错误: {str(e)}")
|
||
|
||
return warnings
|
||
|
||
def _generate_diff(
|
||
self, file_contexts: Dict[str, FileInfo], file: GeneratedFile
|
||
) -> str:
|
||
"""生成 unified diff"""
|
||
if file.is_new:
|
||
return f"--- /dev/null\n+++ {file.path}\n@@ -0,0 +1,{file.content.count(chr(10))} @@\n" + "\n".join(
|
||
f"+{line}" for line in file.content.split("\n")
|
||
)
|
||
|
||
# 对于修改的文件
|
||
old_content = file_contexts.get(file.path)?.content or ""
|
||
old_lines = old_content.split("\n")
|
||
new_lines = file.content.split("\n")
|
||
|
||
# 简化的 diff 生成（实际应该使用 difflib）
|
||
diff_lines = [f"--- {file.path}", f"+++ {file.path}"]
|
||
|
||
# 找出差异的行范围
|
||
changed_ranges = self._find_changed_ranges(old_lines, new_lines)
|
||
|
||
for start_old, start_new, length in changed_ranges:
|
||
diff_lines.append(f"@@ -{start_old + 1},{length} +{start_new + 1},{length} @@")
|
||
for i in range(length):
|
||
old_line = old_lines[start_old + i] if start_old + i < len(old_lines) else ""
|
||
new_line = new_lines[start_new + i] if start_new + i < len(new_lines) else ""
|
||
|
||
if old_line == new_line:
|
||
diff_lines.append(f" {new_line}")
|
||
elif new_line and not old_line:
|
||
diff_lines.append(f"+{new_line}")
|
||
elif old_line and not new_line:
|
||
diff_lines.append(f"-{old_line}")
|
||
else:
|
||
diff_lines.append(f"-{old_line}")
|
||
diff_lines.append(f"+{new_line}")
|
||
|
||
return "\n".join(diff_lines)
|
||
|
||
def _find_changed_ranges(
|
||
self, old_lines: List[str], new_lines: List[str]
|
||
) -> List[tuple]:
|
||
"""找出变化的代码行范围"""
|
||
# 简化实现：假设整个文件是变化的
|
||
return [(0, 0, max(len(old_lines), len(new_lines)))]
|
||
|
||
async def _extract_review_points(
|
||
self, generated_files: List[GeneratedFile],
|
||
file_contexts: Dict[str, FileInfo],
|
||
request: CodeDraftRequest,
|
||
) -> List[str]:
|
||
"""提取需要人工 review 的关键点"""
|
||
review_points = []
|
||
|
||
# 收集所有文件的 review points
|
||
for file in generated_files:
|
||
review_points.extend(file.review_points)
|
||
|
||
# 添加通用 review points
|
||
review_points.extend([
|
||
"验证生成的代码是否符合项目的编码规范",
|
||
"检查是否有安全漏洞（SQL注入、XSS等）",
|
||
"确认复杂业务逻辑的正确性",
|
||
"验证错误处理是否完整",
|
||
"检查是否有性能优化空间",
|
||
])
|
||
|
||
# 根据战略分析添加特定 review points
|
||
strategic_tags = request.strategic_analysis.get("decision_proposal", {}).get("tags", [])
|
||
if "security" in strategic_tags:
|
||
review_points.append("重点关注安全相关代码")
|
||
if "performance" in strategic_tags:
|
||
review_points.append("验证性能指标是否达标")
|
||
|
||
return list(set(review_points)) # 去重
|
||
|
||
def _calculate_confidence(
|
||
self, generated_files: List[GeneratedFile], file_contexts: Dict[str, FileInfo]
|
||
) -> float:
|
||
"""计算生成置信度"""
|
||
if not generated_files:
|
||
return 0.0
|
||
|
||
confidence = 0.7 # 基础置信度
|
||
|
||
# 根据文件数量调整
|
||
confidence -= min(len(generated_files) * 0.05, 0.2)
|
||
|
||
# 根据警告数量调整
|
||
total_warnings = sum(len(f.review_points) for f in generated_files)
|
||
confidence -= min(total_warnings * 0.02, 0.3)
|
||
|
||
return max(0.0, min(1.0, confidence))
|
||
|
||
def _estimate_complexity(self, generated_files: List[GeneratedFile]) -> str:
|
||
"""估算代码复杂度"""
|
||
total_lines = sum(len(f.content.split("\n")) for f in generated_files)
|
||
|
||
if total_lines < 100:
|
||
return "low"
|
||
elif total_lines < 300:
|
||
return "medium"
|
||
else:
|
||
return "high"
|
||
|
||
def _generate_summary(
|
||
self, generated_files: List[GeneratedFile], request: CodeDraftRequest
|
||
) -> str:
|
||
"""生成变更摘要"""
|
||
new_files = [f for f in generated_files if f.is_new]
|
||
modified_files = [f for f in generated_files if not f.is_new]
|
||
|
||
summary_parts = [
|
||
f"为信号 '{request.signal_content[:50]}...' 生成代码草案",
|
||
f"新建文件: {len(new_files)} 个",
|
||
f"修改文件: {len(modified_files)} 个",
|
||
]
|
||
|
||
if modified_files:
|
||
summary_parts.append(f"主要变更: {modified_files[0].path}")
|
||
|
||
return " | ".join(summary_parts)
|
||
|
||
def _collect_warnings(
|
||
self, generated_files: List[GeneratedFile], file_contexts: Dict[str, FileInfo]
|
||
) -> List[str]:
|
||
"""收集警告信息"""
|
||
warnings = []
|
||
|
||
for file in generated_files:
|
||
if not file.content.strip():
|
||
warnings.append(f"{file.path}: 文件内容为空")
|
||
|
||
# 检查是否有 TODO/FIXME
|
||
if "TODO" in file.content or "FIXME" in file.content:
|
||
warnings.append(f"{file.path}: 包含未完成的标记")
|
||
|
||
return warnings
|
||
|
||
def _format_json(self, obj: Any) -> str:
|
||
"""格式化 JSON"""
|
||
import json
|
||
return json.dumps(obj, ensure_ascii=False, indent=2)
|
||
|
||
|
||
# 全局服务实例
|
||
_codegen_service = None
|
||
|
||
|
||
def get_codegen_service(llm_service) -> CodeDraftService:
|
||
"""获取代码生成服务实例"""
|
||
global _codegen_service
|
||
if _codegen_service is None:
|
||
_codegen_service = CodeDraftService(llm_service)
|
||
return _codegen_service
|