Skip to content

CLAUDE.md

Latest commit

 

History

History
284 lines (225 loc) · 9.22 KB

CLAUDE.md

File metadata and controls

284 lines (225 loc) · 9.22 KB

Wake AI Architecture

Overview

Wake AI is now a standalone Python package that provides AI-powered smart contract analysis capabilities. It has been extracted from the Wake framework into an independent module that can be used separately or integrated with other tools.

Module Structure

The project is organized into two main directories:

1. Core Framework (wake_ai/)

Contains the AI framework implementation:

  • core/ - Core Claude integration and workflow engine
    • claude.py - Claude Code CLI wrapper with session management
    • flow.py - Base workflow and step classes
    • exceptions.py - Custom exceptions
    • utils.py - Framework utilities
  • detections.py - Detection data models and formatters
  • results.py - Result types (AIResult, SimpleResult, AIDetectionResult)
  • runner.py - Workflow execution helper
  • utils.py - Shared utilities (YAML loading, validation)
  • cli.py - Command-line interface

2. Workflows (flows/)

Contains pre-built workflow implementations:

  • audit/ - Comprehensive security audit workflow
  • uniswap_detector.py - Uniswap-specific detector

3. Examples (examples/)

Contains educational example workflows:

  • reentrancy/ - Reentrancy detection workflow
  • reentrancy_test/ - Reentrancy testing workflow
  • hooks/ - Hook example workflow

Key Features

Standalone CLI

Wake AI provides its own CLI interface:

wake-ai --flow audit                    # Run audit workflow
wake-ai --flow audit -s contracts/*.sol # Audit specific files
wake-ai --resume                        # Resume previous session
wake-ai --export results.json           # Export results

Session Management

  • Each workflow creates its own ClaudeCodeSession
  • Sessions can be resumed between runs
  • Working directory isolation per session
  • State persistence for interrupted workflows

Cost Management

  • query_with_cost() implements intelligent cost-limited execution
  • Configurable cost limits per step
  • Automatic prompt optimization when approaching limits
  • Cost tracking and reporting

Validation System

  • Step-level validation with automatic retries
  • Custom validators per workflow step
  • Error correction prompts on validation failure

Technical Implementation

Core Classes

AIWorkflow (Base Class)

  • Abstract base for all workflows
  • Manages step execution and state
  • Handles context passing between steps
  • Provides resume capability

WorkflowStep

  • Represents a single workflow step
  • Contains prompt template, tools, and validation
  • Configurable retry and cost limits

ClaudeCodeSession

  • Wrapper around Claude Code CLI
  • Manages session lifecycle
  • Tracks costs and usage
  • Handles working directory setup

Execution Flow

  1. Initialization

    • Workflow creates unique session ID
    • Working directory is created at .wake/ai/<session-id>/
    • Initial context is prepared

  2. Step Execution

    • Each step gets its own ClaudeCodeSession
    • Prompt is rendered with current context
    • Claude executes with specified tools
    • Results are validated (if validator provided)
    • Context is updated with step output

  3. Error Handling

    • Validation failures trigger retries
    • Cost limits are monitored
    • Sessions can be resumed on interruption

Context Management

Wake AI provides methods to manage data flow between steps:

# Add data to context
workflow.add_context("key", value)

# Retrieve data
value = workflow.get_context("key")

# List all keys
keys = workflow.get_context_keys()

Each step has access to:

  • {{working_dir}} - Session working directory
  • {{execution_dir}} - Where workflow was launched
  • Previous step outputs as {{step_name}_output}}
  • User-defined context via add_context()
  • Extracted data from add_extraction_step()

Context is automatically saved/restored when resuming workflows.

Template Syntax

Wake AI uses Jinja2 templating for prompt templates:

  • Variables are referenced with double curly braces: {{variable_name}}
  • Code examples with curly braces no longer need escaping
  • Supports Jinja2 features like conditionals and loops (if needed)
  • Missing variables will raise errors to catch typos early

Working Directory Structure

Each workflow session creates an isolated working directory:

.wake/ai/<session-id>/
├── state/              # Workflow state for resume capability
│   ├── workflow.json   # Workflow metadata and progress
│   └── context.json    # Current context state
├── results/            # AI-generated output files
│   ├── detections.yaml # Security findings
│   ├── report.md       # Analysis reports
│   └── ...             # Other workflow outputs
└── temp/               # Temporary working files

Session Management

  • Session ID Format: YYYYMMDD_HHMMSS_random (e.g., 20250121_143022_abc123)
  • Path Template: .wake/ai/<session-id>/
  • Automatic Creation: Directory created on workflow initialization
  • Context Access: Available as {{working_dir}} in all prompts
  • Automatic Cleanup: Working directories can be automatically cleaned up after successful completion
    • Base AIWorkflow default: cleanup_working_dir = False (preserves working directory)
    • Individual workflows can override this default (e.g., some workflows might set cleanup_working_dir = True)
    • Override via CLI: --no-cleanup to preserve, --cleanup to force cleanup
    • Override in code: pass cleanup_working_dir parameter to workflow constructor

State Persistence

The workflow state is saved after each step:

  • Progress tracking for resume capability
  • Context preservation between steps
  • Cost accumulation tracking
  • Step outputs stored for reference

Creating Custom Workflows

To create a new workflow, inherit from AIWorkflow and create a factory function:

import rich_click as click
from wake_ai import workflow, AIWorkflow, WorkflowStep
from wake_ai.core.utils import validate_yaml_output

@workflow.command("my-workflow")
@click.option("--target", "-t", type=click.Path(exists=True), help="Target file or directory")
def factory(target):
    """Custom workflow implementation."""
    workflow = MyWorkflow()
    workflow.target = target
    return workflow

class MyWorkflow(AIWorkflow):
    """Custom workflow implementation."""

    target: str

    def _setup_steps(self):
        """Define workflow steps."""
        # Step 1: Analysis
        self.add_step(
            name="analyze",
            prompt_template=self._load_prompt("analyze.md"),
            allowed_tools=["Read", "Grep", "Write"],
            max_retries=2,
            max_cost=5.0
        )

        # Step 2: Generate results
        self.add_step(
            name="generate",
            prompt_template=self._load_prompt("generate.md"),
            allowed_tools=["Write"],
            validator=validate_yaml_output,
            max_retries=3
        )

Package Configuration

Installation

# From PyPI
pip install wake-ai

# Development mode
git clone https://github.com/wakehacker/wake-ai
cd wake-ai
uv sync --extra dev

Dependencies

  • Core: click, rich, pyyaml, jinja2, pydantic
  • Development: black, isort, mypy
  • Runtime: Claude Code CLI must be installed

Entry Points

  • wake-ai - Main CLI command
  • Python API: import wake_ai

Migration from Wake

The AI module has been extracted from Wake with the following changes:

  1. Independent Package: No Wake dependencies
  2. Standalone CLI: wake-ai instead of wake ai
  3. Simplified Structure: Focused solely on AI workflows
  4. No Detector Integration: Pure workflow execution
  5. Portable Workflows: Can be used in any project

CLI Registration Changes

# Current approach - factory function pattern
@workflow.command("audit")
@click.option("--scope", "-s", multiple=True, help="Files to audit")
def factory(scope):
    """Run audit workflow."""
    workflow = AuditWorkflow()
    workflow.scope_files = scope
    return workflow

Prompt Writing Guidelines

When creating Wake AI workflow prompts, follow the guidelines in prompt-writing.mdc. Key principles:

Structure

  • Use Task-First Architecture with <task>, <context>, <steps> structure
  • Include explicit <validation_requirements> sections
  • Provide complete <output_format> examples with YAML/code
  • Follow the three-phase pattern: Discovery → Analysis → Documentation

Best Practices

  • Start with a concise one-sentence <task> declaration
  • Use numbered steps with bold headings
  • Include sub-steps (a, b, c) for complex operations
  • Reference specific Wake tools and commands
  • Provide real-world code examples in output formats

Development Notes

TODO Management

The project maintains a TODO list at .claude/TODO.md which tracks:

  • In-progress features and enhancements
  • Prioritized backlog items (High/Medium/Low)
  • Future ideas and recently completed work

Important: When implementing changes or completing features, please update the TODO file to reflect the current state. Mark completed items, add comment if needed.

Design Documentation

Design discussions and specifications are stored in .claude/thoughts/.

You have access to this directory for whenever you need to write down design notes, brainstorming, or other thoughts.