Python Superpowers

AI is great at understanding language. But sometimes you need Python's precision: parsing JSON, filtering data, doing math, or talking to APIs. Code steps let you mix Python into your AI workflows like a secret ingredient.

Your First Code Step

The simplest code step looks like this:

steps:
  - type: code
    code: |
      result = input.upper()
    output_var: result

That's it. Whatever you assign to result becomes available as {result} in subsequent steps.

A Real Example: Smart Word Counter

Let's build something useful—a word counter that also finds the most common words:

name: wordstats
version: "1.0.0"
description: Analyze text statistics

steps:
  - type: code
    code: |
      import re
      from collections import Counter

      # Clean and split
      words = re.findall(r'\b\w+\b', input.lower())

      # Calculate stats
      word_count = len(words)
      unique_count = len(set(words))
      char_count = len(input)

      # Find top words
      top_words = Counter(words).most_common(5)
      top_formatted = '\n'.join(f"  {w}: {c}" for w, c in top_words)
    output_var: word_count, unique_count, char_count, top_formatted

output: |
  Words: {word_count}
  Unique: {unique_count}
  Characters: {char_count}

  Top 5 words:
  {top_formatted}

The Code Step Cookbook

Copy-paste these recipes for common tasks:

- type: code
  code: |
    import json
    try:
        data = json.loads(ai_response)
        result = json.dumps(data, indent=2)
    except json.JSONDecodeError:
        result = ai_response  # Fallback to raw
  output_var: result

- type: code
  code: |
    import re
    # Find all emails
    emails = re.findall(r'[\w.-]+@[\w.-]+\.\w+', input)
    result = '\n'.join(emails) or "No emails found"
  output_var: result

- type: code
  code: |
    lines = [l for l in input.split('\n') if l.strip()]
    cleaned = '\n'.join(lines)
  output_var: cleaned

- type: code
  code: |
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    result = f"[{timestamp}] {input}"
  output_var: result

- type: code
  code: |
    max_len = int(max_length)  # From --max-length flag
    if len(input) > max_len:
        truncated = input[:max_len] + "..."
    else:
        truncated = input
  output_var: truncated

Bulletproof Error Handling

- type: code
  code: |
    import json

    try:
        # Risky operation
        data = json.loads(ai_response)
        items = data.get('items', [])
        result = '\n'.join(items)
        error = None
    except json.JSONDecodeError as e:
        # Graceful fallback
        result = ai_response
        error = f"Parse warning: {e}"
    except KeyError as e:
        result = "Missing expected field"
        error = str(e)
  output_var: result, error

The AI + Code Power Combo

The real magic happens when you alternate between AI and code:

name: structured-extract
version: "1.0.0"
description: Extract and validate structured data

steps:
  # AI extracts data
  - type: prompt
    provider: claude
    prompt: |
      Extract all dates from this text.
      Return as JSON array: ["YYYY-MM-DD", ...]

      {input}
    output_var: dates_json

  # Code validates and sorts
  - type: code
    code: |
      import json
      from datetime import datetime

      try:
          dates = json.loads(dates_json)
          # Validate each date
          valid = []
          for d in dates:
              try:
                  datetime.strptime(d, "%Y-%m-%d")
                  valid.append(d)
              except ValueError:
                  pass
          # Sort chronologically
          sorted_dates = sorted(valid)
          result = '\n'.join(sorted_dates)
          count = len(sorted_dates)
      except json.JSONDecodeError:
          result = "Could not parse dates"
          count = 0
    output_var: result, count

  # AI formats nicely
  - type: prompt
    provider: claude
    prompt: |
      Format these {count} dates in a human-readable way:
      {result}
    output_var: formatted

output: "{formatted}"

Try It: Build a CSV Analyzer

name: csv-describe
version: "1.0.0"
description: Understand CSV data at a glance

steps:
  - type: code
    code: |
      import csv
      from io import StringIO

      reader = csv.reader(StringIO(input))
      rows = list(reader)

      if rows:
          headers = rows[0]
          row_count = len(rows) - 1
          sample = rows[1:4]  # First 3 data rows
          headers_str = ', '.join(headers)
          sample_str = '\n'.join(', '.join(r) for r in sample)
      else:
          headers_str = "No headers"
          row_count = 0
          sample_str = "No data"
    output_var: headers_str, row_count, sample_str

  - type: prompt
    provider: claude
    prompt: |
      This CSV has {row_count} rows with these columns:
      {headers_str}

      Sample data:
      {sample_str}

      What does this data likely represent?
      What insights could we extract from it?
    output_var: analysis

output: |
  Columns: {headers_str}
  Rows: {row_count}

  {analysis}

What You Can Import

The Python standard library is fully available:

Third-party packages work too—just make sure they're installed in your environment.

Security: The Rules

Code steps run with your user permissions. Treat input as untrusted data—parse it, don't execute it.

Level Up

Ready for the advanced stuff?

• Advanced Workflows - Multi-provider, conditional logic, external tools
• Parallel Orchestration - Run multiple tools simultaneously