advanced_example.py

"""Advanced example with dependencies and complex tools."""

from dataclasses import dataclass
from typing import List
from pydantic_ai import Agent, RunContext
from claude_codemode import codemode, CodeModeConfig


# Define dependencies
@dataclass
class DatabaseContext:
    """Simulated database context."""

    connection_string: str = "sqlite:///:memory:"
    max_retries: int = 3


@dataclass
class APIContext:
    """Simulated API context."""

    api_key: str = "sk-test-key"
    base_url: str = "https://api.example.com"


# Create agent with dependencies
agent = Agent(
    'claude-sonnet-4-5-20250929',
    deps_type=tuple[DatabaseContext, APIContext]
)


@agent.tool
def query_database(ctx: RunContext[tuple[DatabaseContext, APIContext]], query: str) -> List[dict]:
    """Execute a database query.

    Args:
        ctx: Runtime context with database connection
        query: SQL query to execute

    Returns:
        List of query results
    """
    db_ctx, _ = ctx.deps

    # Simulated database query
    mock_data = {
        "SELECT * FROM users": [
            {"id": 1, "name": "Alice", "email": "alice@example.com"},
            {"id": 2, "name": "Bob", "email": "bob@example.com"},
        ],
        "SELECT * FROM orders": [
            {"id": 101, "user_id": 1, "total": 150.00},
            {"id": 102, "user_id": 2, "total": 200.00},
        ],
    }

    return mock_data.get(query, [])


@agent.tool
def call_external_api(
    ctx: RunContext[tuple[DatabaseContext, APIContext]], endpoint: str, params: dict
) -> dict:
    """Call an external API.

    Args:
        ctx: Runtime context with API credentials
        endpoint: API endpoint to call
        params: Request parameters

    Returns:
        API response
    """
    _, api_ctx = ctx.deps

    # Simulated API call
    mock_responses = {
        "/enrichment": {
            "data": {
                "alice@example.com": {"company": "Acme Corp", "title": "Engineer"},
                "bob@example.com": {"company": "Tech Inc", "title": "Manager"},
            }
        },
        "/analytics": {
            "metrics": {
                "total_revenue": 350.00,
                "avg_order_value": 175.00,
                "customer_count": 2,
            }
        },
    }

    return mock_responses.get(endpoint, {"error": "Endpoint not found"})


@agent.tool
def process_data(data: List[dict], operation: str) -> dict:
    """Process data with various operations.

    Args:
        data: List of data records
        operation: Operation to perform (sum, count, filter, etc.)

    Returns:
        Processed result
    """
    if operation == "sum":
        # Sum numeric fields
        total = sum(
            sum(v for v in record.values() if isinstance(v, (int, float)))
            for record in data
        )
        return {"operation": "sum", "result": total, "count": len(data)}

    elif operation == "count":
        return {"operation": "count", "result": len(data)}

    elif operation == "aggregate":
        # Aggregate by type
        numeric_sum = 0
        string_count = 0
        for record in data:
            for value in record.values():
                if isinstance(value, (int, float)):
                    numeric_sum += value
                elif isinstance(value, str):
                    string_count += 1

        return {
            "operation": "aggregate",
            "numeric_sum": numeric_sum,
            "string_count": string_count,
            "total_records": len(data),
        }

    return {"operation": operation, "error": "Unknown operation"}


@agent.tool
def generate_report(title: str, sections: List[dict]) -> str:
    """Generate a formatted report.

    Args:
        title: Report title
        sections: List of report sections with 'name' and 'content' keys

    Returns:
        Formatted report string
    """
    report = f"# {title}\n\n"

    for section in sections:
        name = section.get("name", "Untitled")
        content = section.get("content", "")

        report += f"## {name}\n\n"
        report += f"{content}\n\n"

    report += "---\n"
    report += f"Generated by Claude Codemode\n"

    return report


def main():
    """Run an advanced codemode example."""
    print("=" * 80)
    print("Advanced Claude Codemode Example")
    print("=" * 80)

    # Create dependencies
    db_ctx = DatabaseContext(connection_string="sqlite:///app.db")
    api_ctx = APIContext(api_key="sk-prod-key-12345")
    deps = (db_ctx, api_ctx)

    # Configure codemode
    config = CodeModeConfig(
        verbose=True,
        preserve_workspace=True,
        timeout=120,
    )

    # Complex multi-step task
    prompt = """
    Generate a comprehensive customer analytics report with the following steps:

    1. Query the database for all users and orders
    2. Enrich user data by calling the /enrichment API endpoint
    3. Get analytics metrics from the /analytics API endpoint
    4. Process the orders data to calculate aggregated statistics
    5. Generate a formatted report with the following sections:
       - Customer Overview (user data with enrichment)
       - Order Analytics (metrics from API)
       - Statistical Summary (processed aggregates)

    Return the final formatted report.
    """

    print(f"\nPrompt: {prompt}\n")

    result = codemode(agent, prompt, deps=deps, config=config)

    print("\n" + "=" * 80)
    print("Result")
    print("=" * 80)

    if result.success:
        print(f"✓ Success!\n")
        print(result.output)
    else:
        print(f"✗ Failed\n")
        print(f"Error: {result.error}")

    if config.verbose:
        print(f"\n{'-' * 80}")
        print("Execution log:")
        print(f"{'-' * 80}")
        print(result.execution_log)


if __name__ == "__main__":
    main()