wakehacker/wake-ai

# Wake AI 一个 LLM 编排框架，用于封装终端式 AI 代理（如 Claude Code），提供结构化的多步骤工作流，用于智能合约安全分析及其他场景，使代理执行更可预测、更可靠，通过验证和渐进式任务分解实现。  ## 问题 传统 AI 驱动的代码分析方法存在不可预测性。你编写提示词，交叉手指，希望 AI 一次性正确完成所有工作。当涉及复杂的安全审计或漏洞检测时，这种方法存在根本缺陷。AI 代理功能强大但本质上是不确定的——它们可能会遗漏步骤、产生不一致的输出，或在执行过程中失败。 ## 我们的解决方案 Wake AI 在 AI 的泛化能力与可靠、可重现结果的需求之间搭建桥梁。通过将复杂任务分解为具有中间验证的离散步骤，我们实现了： - **可预测的执行**：每个步骤都有明确的输入、输出和成功标准 - **渐进式验证**：在继续之前验证 AI 是否已完成必要工作 - **多代理协作**：不同步骤可使用专用代理 - **快速原型**：无需从头构建即可测试新想法 ## 安装 ``` pip install wake-ai ``` ## 快速开始 ``` # 运行全面安全审计 wake-ai audit # 检测重入漏洞 wake-ai reentrancy ``` ## 为何选择 Wake AI？ ### 1. **超越静态分析** 我们正进入漏洞检测的新时代。传统检测器依赖模式匹配和静态分析，而 Wake AI 利用 LLM 理解上下文、推理代码行为，并发现基于规则的系统遗漏的复杂漏洞。 ### 2. **结构化且灵活** 该框架在不牺牲灵活性的情况下提供结构： - 将工作流定义为一系列步骤 - 每个步骤可使用不同的工具和提示词 - 验证确保质量后再继续 - 上下文在各步骤间无缝流动 ### 3. **可控成本执行** 现实世界的 AI 使用需要成本管理： - 设置每步成本上限 - 失败时自动重试并反馈 - 接近限制时优化提示词 - 跟踪整个工作流的总成本 ### 4. **进度追踪** 长时间运行工作流的视觉反馈： - 带有完成百分比的进度条 - 重试和验证期间的状态信息 - 用于应用集成的外部进度钩子 ### 5. **内置安全工作流** - 提供用于入门的预构建审计和检测器工作流 - 带 JSON 导出的标准化检测输出格式 - 可构建任何你需要的工作流——安全、测试、分析或其他 - 易于使用自定义工作流扩展 ## 核心概念 ### 工作目录 多步骤 AI 工作流面临一个基本问题：上下文共享。无论你是在所有步骤中使用单个代理还是多个专用代理，数据都需要在操作间流动。漏洞开发代理如何访问审计代理发现的漏洞？第三步如何知道第一步发现了什么？传统方法需要复杂的状太管理或强制 AI 重新分析所有内容。 Wake AI 采用直接方法：每个工作流获得一个隔离的工作目录，所有代理和步骤在此操作。这个共享工作空间成为工作流的持久记忆： ``` .wake/ai// ├── state/ # Workflow state metadata ├── findings.yaml # Discovered vulnerabilities ├── analysis.md # Detailed investigation notes └── report.json # Final structured output ``` 通信通过文件进行。一个代理写入发现，另一个读取并验证。安全审计员记录漏洞，漏洞开发代理测试它们，报告撰写者整合所有内容。无状态传递，无变量 juggling——只有一个在完整工作流执行期间持续的共享文件系统。你的项目目录保持干净，而每个代理可自由创建、修改和引用此沙箱中的文件。 #### 工作流后清理 默认情况下，工作流配置为在成功完成后自动清理其工作目录。此行为可在 `Workflow` 类中针对特定工作流覆盖，也可在命令行中覆盖： ``` wake-ai --working-dir .audit/ --no-cleanup audit # Preserve working directory ``` 某些工作流可能不需要结构化输出，而是直接在工作目录中提供结果。一个示例是专门的 `audit` 工作流，其输出写入工作目录中的 Markdown 文件，供人工审计员在工作流结束后审查。 ### 验证 为确保整个工作流的正确输出，每个步骤必须产生正确的中间结果。缺乏验证时，AI 响应可能不可解析、缺少必要字段，或包含会传播到后续步骤的错误。 每个工作流步骤可包含验证器，以确保在继续之前输出正确。验证失败时，步骤会自动重试并使用错误纠正提示，直到输出满足要求。 示例验证器： ``` def validate_findings(self, response): # Check if required files were created if not (self.working_dir / "findings.yaml").exists(): return False, ["No findings file created"] # Validate YAML structure findings = yaml.load(open(self.working_dir / "findings.yaml")) if not all(k in findings for k in ["vulnerabilities", "severity"]): return False, ["Invalid findings structure"] return True, [] ``` **优点：** - **自纠正**：AI 自动修复验证错误 - **质量控制**：输出始终符合规范 - **无级联错误**：无效输出不会传播到后续步骤 - **成本高效**：如果验证失败，长时间运行的工作流可提前终止 - **基于架构的输出**：验证器强制特定输出格式，使 AI 响应可解析为结构化数据以便导出并与其他工具集成 ## 架构 ### 工作流与步骤 核心上，Wake AI 将复杂任务视为由独立步骤组成的工作流： ``` from wake_ai import AIWorkflow class MyAuditWorkflow(AIWorkflow): def _setup_steps(self): # Step 1: Map the codebase structure self.add_step( name="map_contracts", prompt_template="Find all Solidity contracts and identify their relationships...", validator=self.validate_mapping, max_cost=5.0 ) # Step 2: Focus on critical contracts (continues session) self.add_step( name="analyze_critical", prompt_template="Based on the contracts you just mapped, analyze the 3 most critical ones for vulnerabilities...", max_cost=10.0, continue_session=True # Needs to remember which contracts were identified ) ``` **特性：** - **成本控制**：若设置 `max_cost`，达到上限后代理将被提示快速完成步骤，对管理高成本步骤有用 - **会话延续**：`continue_session` 控制是否从上一延续代理会话或创建新会话，允许在单代理或多代理工作流间选择 - **验证器**：每个步骤可拥有验证函数，用于检查步骤是否成功完成 - **条件执行**：`condition` 可用于基于布尔表达式跳过步骤 **注意**：步骤按顺序执行。Wake AI 当前不支持并行步骤执行。 ### 上下文共享 Wake AI 上下文共享的主要方法是通过 _工作目录_。此外，工作流还提供上下文管理方法和 `add_extraction_step` 辅助函数以提取结构化数据。 ``` from pydantic import BaseModel class Vulnerability(BaseModel): name: str description: str severity: str file: str line: int class VulnerabilitiesList(BaseModel): vulnerabilities: List[Vulnerability] self.add_step( name="analyze_critical", prompt_template="Analyze the codebase for critical vulnerabilities...", ) self.add_extraction_step( after_step="analyze_critical", output_schema=VulnerabilitiesList, context_key="vulnerabilities", ) ``` 提取的数据将存储在 `context` 状态中，键名为 `context_key`（默认为 `_data`）。 ### 上下文管理 Wake AI 工作流提供管理跨步骤流动上下文的方法： ``` # 添加上下文数据 workflow.add_context("project_name", "MyDeFiProtocol") workflow.add_context("audit_scope", ["Token.sol", "Vault.sol"]) # 检索上下文数据 project = workflow.get_context("project_name") # "MyDeFiProtocol" scope = workflow.get_context("audit_scope") # ["Token.sol", "Vault.sol"] # 获取所有可用的上下文键 keys = workflow.get_context_keys() # ["project_name", "audit_scope", ...] ``` 上下文包括： - 通过 `add_context()` 添加的用户定义值 - 步骤输出 `{{step_name}_output}}` - 通过 `add_extraction_step()` 提取的数据（默认为 `_data`，但可通过 `context_key` 参数覆盖） - 内置值如 `{{working_dir}}` 和 `{{execution_dir}}` ### 动态提示模板 为创建动态提示模板，Wake AI 使用 Jinja2 模板，允许以 `{{ context_key }}` 格式传递上下文变量。工作流类跟踪其 `context` 状态，可在其中存储任意希望传递到提示模板的数据。 ``` from pydantic import BaseModel class ContractList(BaseModel): contracts: str class AuditWorkflow(AIWorkflow): def __init__(self, **kwargs): super().__init__(**kwargs) # Add files to context for use in prompts self.add_context("files", list(self.working_dir.glob("**/*.sol"))) def _setup_steps(self): self.add_step( name="map_contracts", prompt_template="Find all Solidity contracts and identify their relationships. Here are the files in the codebase: {{files}}", ) self.add_step( name="determine_focus", prompt_template="Determine the top 3 core contracts to focus on for analysis.", continue_session=True # Needs to remember which contracts were identified ) self.add_extraction_step( after_step="determine_focus", output_schema=ContractList, context_key="files_to_focus", ) self.add_step( name="analyze_focus", prompt_template="Conduct a thorough analysis of the following contracts: {{files_to_focus}}", max_cost=10.0, ) ``` ## 创建自定义检测器 Wake AI 使使用 `SimpleDetector` 辅助类快速构建新的漏洞检测器变得非常简单： ``` import rich_click as click from wake_ai import workflow from wake_ai.templates import SimpleDetector @workflow.command("flashloan") @click.option("--focus-area", "-f", help="Specific area to focus analysis") def factory(focus_area): """Detect flash loan attack vectors.""" detector = FlashLoanDetector() detector.focus_area = focus_area return detector class FlashLoanDetector(SimpleDetector): """Flash loan vulnerability detector.""" focus_area: str = None def get_detector_prompt(self) -> str: focus_context = f"Focus specifically on: {self.focus_area}" if self.focus_area else "" return f""" Analyze this codebase for flash loan attack vectors. {focus_context} Focus on: 1. Price manipulation opportunities 2. Unprotected external calls in loan callbacks 3. Missing reentrancy guards 4. Incorrect balance assumptions For each issue, explain the attack scenario and impact. """ ``` 该辅助类会自动将检测器输出解析为标准格式，用于 CLI 显示或 JSON 导出。 ## 高级特性 ### 动态步骤生成 基于运行时发现生成步骤： ``` def generate_file_reviews(response, context): # Parse discovered contracts contracts = parse_contracts(response.content) # Create a review step for each contract return [ WorkflowStep( name=f"review_{contract.name}", prompt_template=f"Audit {contract.path} for vulnerabilities...", max_cost=3.0 ) for contract in contracts ] self.add_dynamic_steps("reviews", generate_file_reviews, after_step="scan") ``` ### 条件执行 基于运行时条件跳过昂贵步骤： ``` self.add_step( name="deep_analysis", prompt_template="Perform computational analysis...", condition=lambda ctx: len(ctx.get("critical_findings", [])) > 0, max_cost=20.0 ) ``` ### 工具限制 对 AI 能力进行细粒度控制： ``` # 仅允许特定 Wake 命令 allowed_tools = ["Read", "Write", "Bash(wake detect *)", "Bash(wake init)"] # 防止任何修改 self.add_step( name="review_only", prompt_template="Review the code...", disallowed_tools=["Write", "Edit", "Bash"] ) ``` ## 实际安全审计工作流示例 我们的生产审计工作流展示了框架的能力： 1. **初始分析**——映射攻击面并确定重点区域 2. **漏洞探测**——使用专用提示词深入挖掘 3. **报告生成**——生成专业的审计文档 每个步骤都验证输出，确保不会遗漏关键检查。 ## 文档 请参阅 [docs/README.md](docs/README.md) 获取： - 完整的 API 参考 - 分步教程 - 高级工作流模式 - 集成指南 ## 许可证 ISC

标签：AI编排框架, LLM编排, pip安装, Python包, Reentrancy检测, SEO: AI安全分析, SEO: 可预测AI执行, SEO: 智能合约审计, 上下文理解, 云安全监控, 区块链安全, 可靠性, 可预测执行, 多代理协作, 多步骤工作流, 审计工具, 工作流分解, 快速原型, 智能合约安全分析, 杀软绕过, 渐进式验证, 终端AI代理, 结构化流程, 输入输出验证, 逆向工具, 静态分析