代码审计智能体设计文档

from typing import TypedDict, Annotated, List, Dict, Optional
import operator
from langchain_core.messages import BaseMessage

class CodeAuditState(TypedDict):
    """代码审计智能体状态 - 使用TypedDict确保类型安全"""

    # === 输入 ===
    repository: str
    target_user: str
    branch: Optional[str]
    commits_limit: int
    audit_dimensions: List[str]

    # === 消息历史(支持自动追加) ===
    messages: Annotated[List[BaseMessage], operator.add]

    # === 中间状态 ===
    code_changes: Dict[str, any]
    user_profile: Optional[Dict]
    static_analysis_results: Dict[str, any]
    security_scan_results: Dict[str, any]
    cache_key: Optional[str]
    cache_hit: bool

    # === 审计结果(支持自动追加) ===
    code_quality_issues: Annotated[List[Dict], operator.add]
    security_issues: Annotated[List[Dict], operator.add]
    performance_issues: Annotated[List[Dict], operator.add]
    architecture_issues: Annotated[List[Dict], operator.add]
    code_smells: Annotated[List[Dict], operator.add]
    best_practices: Annotated[List[Dict], operator.add]

    # === 输出 ===
    audit_report: str
    severity_summary: Dict[str, int]
    personalized_suggestions: List[str]
    estimated_cost: float

    # === 流程控制 ===
    next_step: Optional[str]
    has_critical_issues: bool
    error: Optional[str]

from langchain_anthropic import ChatAnthropic
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
import asyncio

# === 节点1: 命令解析 ===
def parse_command(state: CodeAuditState) -> CodeAuditState:
    """解析GitHub Issue中的审计命令"""
    # 从消息中提取命令参数
    # 实现略...
    return {
        "messages": [{"role": "system", "content": "命令解析完成"}],
        "next_step": "fetch_code"
    }

# === 节点2: 代码获取 ===
async def fetch_code_changes(state: CodeAuditState) -> CodeAuditState:
    """从GitHub API获取代码变更"""
    from github import Github

    # 检查缓存
    cache_key = f"{state['repository']}:{state['target_user']}:{state['branch']}"
    cached_result = await check_cache(cache_key)

    if cached_result:
        return {
            "cache_hit": True,
            "cache_key": cache_key,
            "next_step": "generate_report"  # 跳过分析,直接生成报告
        }

    # 获取代码diff
    g = Github(GITHUB_TOKEN)
    repo = g.get_repo(state["repository"])
    commits = repo.get_commits(
        author=state["target_user"],
        since=...,
        until=...
    )[:state["commits_limit"]]

    code_changes = {}
    for commit in commits:
        for file in commit.files:
            code_changes[file.filename] = {
                "patch": file.patch,
                "additions": file.additions,
                "deletions": file.deletions
            }

    return {
        "code_changes": code_changes,
        "cache_hit": False,
        "cache_key": cache_key,
        "messages": [{"role": "system", "content": f"获取到{len(code_changes)}个文件的变更"}],
        "next_step": "parallel_analysis"
    }

# === 节点3-5: 并行分析节点 ===
async def build_user_profile(state: CodeAuditState) -> CodeAuditState:
    """构建用户画像(从历史审计记录学习)"""
    # 查询PostgreSQL历史数据
    user_history = await db.get_user_audit_history(state["target_user"])

    # 分析用户常见问题模式
    common_issues = analyze_common_patterns(user_history)

    return {
        "user_profile": {
            "common_issues": common_issues,
            "tech_stack": ["Java", "Spring"],  # 从历史代码推断
            "improvement_trend": "positive"  # 代码质量趋势
        }
    }

async def static_analysis(state: CodeAuditState) -> CodeAuditState:
    """SonarQube静态分析"""
    # 调用SonarQube API
    issues = await sonarqube_client.analyze(
        files=state["code_changes"].keys()
    )

    return {
        "static_analysis_results": issues,
        "code_quality_issues": [
            {"type": "naming", "severity": "INFO", "file": "...", "line": 42}
        ]
    }

async def security_scan(state: CodeAuditState) -> CodeAuditState:
    """安全漏洞扫描"""
    # 调用OWASP工具
    vulnerabilities = await owasp_scanner.scan(
        code_changes=state["code_changes"]
    )

    return {
        "security_scan_results": vulnerabilities,
        "security_issues": [
            {"type": "sql_injection", "severity": "CRITICAL", "file": "...", "line": 100}
        ]
    }

# === 节点6: AI深度审计 ===
async def ai_deep_audit(state: CodeAuditState) -> CodeAuditState:
    """Claude AI深度代码审计"""
    llm = ChatAnthropic(model="claude-3-5-sonnet-20241022", temperature=0)

    # 构建Prompt
    prompt = f"""
你是资深代码审计专家。请审计以下代码变更:

**代码变更**:
{format_code_changes(state["code_changes"])}

**静态分析结果**:
{state["static_analysis_results"]}

**用户画像**:
{state["user_profile"]}

**审计维度**: {', '.join(state["audit_dimensions"])}

请从以下维度审计代码:
1. 架构设计合理性
2. 性能优化建议
3. 代码异味识别
4. 最佳实践应用

输出JSON格式:
{{
  "architecture_issues": [...],
  "performance_issues": [...],
  "code_smells": [...],
  "best_practices": [...]
}}
"""

    response = await llm.ainvoke(prompt)
    ai_results = parse_json(response.content)

    return {
        "architecture_issues": ai_results["architecture_issues"],
        "performance_issues": ai_results["performance_issues"],
        "code_smells": ai_results["code_smells"],
        "best_practices": ai_results["best_practices"],
        "messages": [{"role": "assistant", "content": "AI审计完成"}]
    }

# === 节点7: 个性化建议 ===
async def personalized_advice(state: CodeAuditState) -> CodeAuditState:
    """基于用户画像生成个性化建议"""
    llm = ChatAnthropic(model="claude-3-5-sonnet-20241022")

    # 汇总所有问题
    all_issues = (
        state["code_quality_issues"] +
        state["security_issues"] +
        state["performance_issues"] +
        state["architecture_issues"] +
        state["code_smells"]
    )

    prompt = f"""
基于用户 {state["target_user"]} 的历史表现和本次审计结果,生成个性化建议:

**用户画像**: {state["user_profile"]}
**本次问题**: {len(all_issues)} 个问题
**问题分布**: {categorize_issues(all_issues)}

请生成:
1. 针对性改进建议(基于用户常见问题模式)
2. 学习资源推荐
3. 进步反馈(对比历史趋势)
"""

    response = await llm.ainvoke(prompt)

    return {
        "personalized_suggestions": parse_suggestions(response.content)
    }

# === 节点8: 报告生成 ===
def generate_report(state: CodeAuditState) -> CodeAuditState:
    """生成Markdown格式审计报告"""

    # 统计严重级别
    severity_summary = calculate_severity_summary(state)

    # 生成Markdown
    report = f"""
# 代码审计报告

**审计对象**: @{state["target_user"]}
**仓库**: {state["repository"]}
**分支**: {state["branch"] or "全部"}
**审计时间**: {datetime.now()}

---

## 📊 问题统计

| 严重级别 | 数量 |
|---------|------|
| CRITICAL | {severity_summary.get('CRITICAL', 0)} |
| HIGH     | {severity_summary.get('HIGH', 0)} |
| MEDIUM   | {severity_summary.get('MEDIUM', 0)} |
| LOW      | {severity_summary.get('LOW', 0)} |
| INFO     | {severity_summary.get('INFO', 0)} |

**总计**: {sum(severity_summary.values())} 个问题

---

## 🔴 严重问题 (CRITICAL)

{format_issues(state["security_issues"], "CRITICAL")}

## 🔵 高优先级问题 (HIGH)

{format_issues(state["performance_issues"], "HIGH")}

## 🟡 中等问题 (MEDIUM)

{format_issues(state["architecture_issues"], "MEDIUM")}

## 🟢 建议改进 (LOW)

{format_issues(state["code_smells"], "LOW")}

---

## 💡 个性化建议

{chr(10).join(f"- {s}" for s in state["personalized_suggestions"])}

---

## 📈 进步趋势

{generate_trend_analysis(state["user_profile"])}

---

**审计成本**: ${state["estimated_cost"]:.3f}
**缓存命中**: {"是" if state["cache_hit"] else "否"}
"""

    # 判断是否有严重问题
    has_critical = severity_summary.get("CRITICAL", 0) > 0

    return {
        "audit_report": report,
        "severity_summary": severity_summary,
        "has_critical_issues": has_critical,
        "next_step": "human_review" if has_critical else "post_to_github"
    }

# === 节点9: GitHub反馈 ===
async def post_to_github(state: CodeAuditState) -> CodeAuditState:
    """将审计报告发布到GitHub Issue"""
    from github import Github

    g = Github(GITHUB_TOKEN)
    repo = g.get_repo(state["repository"])
    issue = repo.get_issue(state["issue_id"])

    # 发布评论
    comment = issue.create_comment(state["audit_report"])

    # 严重问题自动Request Changes
    if state["has_critical_issues"]:
        # PR场景下请求变更
        if hasattr(issue, 'create_review'):
            issue.create_review(
                body="发现严重安全问题,请修复后再合并",
                event="REQUEST_CHANGES"
            )

    # 保存到缓存
    await save_to_cache(
        key=state["cache_key"],
        value=state["audit_report"],
        ttl=86400  # 24小时
    )

    return {
        "messages": [{"role": "system", "content": f"已发布到GitHub: {comment.html_url}"}]
    }

from typing import Literal

def should_use_cache(state: CodeAuditState) -> Literal["parallel_analysis", "generate_report"]:
    """路由: 是否使用缓存结果"""
    if state["cache_hit"]:
        return "generate_report"  # 跳过分析
    return "parallel_analysis"  # 执行分析

def should_human_review(state: CodeAuditState) -> Literal["human_review", "post_to_github"]:
    """路由: 是否需要人工审核"""
    if state["has_critical_issues"]:
        return "human_review"  # 严重问题需人工确认
    return "post_to_github"  # 直接发布

def handle_error(state: CodeAuditState) -> Literal["retry", "fallback", END]:
    """错误处理路由"""
    if state["error"]:
        if state["retry_count"] < 3:
            return "retry"  # 重试
        return "fallback"  # 降级(仅静态分析)
    return END

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

# === 构建StateGraph ===
builder = StateGraph(CodeAuditState)

# === 添加节点 ===
builder.add_node("parse_command", parse_command)
builder.add_node("fetch_code_changes", fetch_code_changes)

# 并行分析节点
builder.add_node("build_user_profile", build_user_profile)
builder.add_node("static_analysis", static_analysis)
builder.add_node("security_scan", security_scan)

# 后续节点
builder.add_node("ai_deep_audit", ai_deep_audit)
builder.add_node("personalized_advice", personalized_advice)
builder.add_node("generate_report", generate_report)
builder.add_node("post_to_github", post_to_github)

# === 添加边 ===
builder.add_edge(START, "parse_command")
builder.add_edge("parse_command", "fetch_code_changes")

# 条件边: 缓存检查
builder.add_conditional_edges(
    "fetch_code_changes",
    should_use_cache,
    {
        "parallel_analysis": "build_user_profile",  # 未命中缓存
        "generate_report": "generate_report"  # 命中缓存,跳过分析
    }
)

# 并行执行三个分析节点(Fan-out模式)
builder.add_edge("build_user_profile", "ai_deep_audit")
builder.add_edge("static_analysis", "ai_deep_audit")
builder.add_edge("security_scan", "ai_deep_audit")

# 注意: 需要在fetch_code_changes节点中同时触发这3个节点
# 或者使用Send API(LangGraph 0.2+新特性)实现真正的并行

# AI审计后续流程
builder.add_edge("ai_deep_audit", "personalized_advice")
builder.add_edge("personalized_advice", "generate_report")

# 条件边: 人工审核
builder.add_conditional_edges(
    "generate_report",
    should_human_review,
    {
        "human_review": "human_review",  # 等待人工确认
        "post_to_github": "post_to_github"  # 直接发布
    }
)

builder.add_edge("post_to_github", END)

# === 编译图(带检查点) ===
memory = MemorySaver()  # 使用内存存储检查点
graph = builder.compile(
    checkpointer=memory,  # 启用状态持久化
    interrupt_before=["human_review"]  # 严重问题时暂停,等待人工确认
)

# === 可视化图结构 ===
from IPython.display import Image, display
display(Image(graph.get_graph().draw_mermaid_png()))

# === 执行审计 ===
initial_state = {
    "repository": "maiban-health/backend",
    "target_user": "zhangsan",
    "branch": "feature/payment-service",
    "commits_limit": 5,
    "audit_dimensions": ["security", "performance", "code_quality"],
    "messages": [],
    "code_quality_issues": [],
    "security_issues": [],
    "performance_issues": [],
    "architecture_issues": [],
    "code_smells": [],
    "best_practices": [],
    "cache_hit": False,
    "has_critical_issues": False,
    "estimated_cost": 0.0,
    "retry_count": 0
}

# 配置(支持断点续传)
config = {"configurable": {"thread_id": "audit-123"}}

# 执行图
result = await graph.ainvoke(initial_state, config)

print(result["audit_report"])
print(f"成本: ${result['estimated_cost']:.3f}")

# === 如果遇到human_review中断 ===
# 1. 查看当前状态
state_snapshot = graph.get_state(config)
print("等待人工审核:", state_snapshot.values["audit_report"])

# 2. 人工确认后继续
# 批准发布
graph.update_state(config, {"next_step": "post_to_github"})
result = await graph.ainvoke(None, config)

# 或者拒绝(修改报告)
graph.update_state(config, {"audit_report": "修改后的报告..."})
result = await graph.ainvoke(None, config)