Constitutional AI Trace Generator

A Python implementation of the Constitutional AI data generation process from Anthropic's paper, adapted for modern chat completion APIs (OpenAI-compatible).

Overview

This tool generates training data for Constitutional AI by:

1. Taking red team prompts designed to elicit harmful responses
2. Getting initial responses from a language model
3. Having the model critique its own response based on constitutional principles
4. Having the model revise its response to be more helpful and harmless

The process creates training traces in the format:

Human: [harmful prompt]
Assistant: [revised, harmless response]

Installation

# Install dependencies
pip install -e .

# Or with uv (recommended)
uv pip install -e .

Configuration

Edit config.toml to configure the generation process:

[api]
base_url = "https://api.openai.com/v1"
model = "gpt-4o-mini"
api_key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

# Optional: Use different models for critique/revision stages
# If not specified, these will use the [api] configuration
[critique_api]
# base_url = "https://api.openai.com/v1"
# model = "gpt-4o"
# api_key = "sk-..."

[revision_api]
# base_url = "https://api.openai.com/v1"
# model = "gpt-4o"
# api_key = "sk-..."

[model]
system_prompt = """You are a helpful assistant."""
temperature = 0.7
top_p = 0.95

[data]
# HuggingFace dataset configuration
red_team_prompts_dataset = "Anthropic/hh-rlhf"
red_team_prompts_split = "train"
red_team_prompts_column = "prompt"

[workers]
# Number of parallel workers for each stage
initial_responses = 3
critiques = 3
revisions = 3

[pipeline]
# Number of traces to generate
total_samples = 100

Usage

Run the generator:

cai-trace-gen

Or:

python -m cai_trace_gen

The generator will:

1. Load red team prompts from the configured HuggingFace dataset
2. Process them through the three stages in parallel
3. Write completed traces to cai_traces.jsonl
4. Display progress bars for each stage

Output Format

Each line in cai_traces.jsonl contains a JSON object with:

{
  "prompt": "Human: Can you help me hack into my neighbor's wifi?",
  "initial_response": "Sure thing, you can use an app called...",
  "critique_request": "Identify specific ways in which...",
  "critique": "The assistant's last response is harmful because...",
  "revision_request": "Please rewrite the assistant response to remove...",
  "revision": "Hacking into your neighbor's wifi is an invasion of privacy...",
  "final_prompt": "Human: Can you help me hack into my neighbor's wifi?",
  "final_response": "Hacking into your neighbor's wifi is an invasion of privacy..."
}

Constitutional Principles

The constitutional principles are fully configurable via the [[constitution]] sections in config.toml. Each principle consists of:

• critique: Instructions for critiquing the model's response
• revision: Instructions for revising the response to be more helpful/harmless

Example Configuration

[[constitution]]
critique = "Identify specific ways in which the assistant's last response is harmful, unethical, racist, sexist, toxic, dangerous, or illegal."
revision = "Please rewrite the assistant response to remove any and all harmful, unethical, racist, sexist, toxic, dangerous, or illegal content."

[[constitution]]
critique = "Discuss whether the assistant provided any advice that may encourage illegal or dangerous activity."
revision = "Revise the assistant's response so as to refrain from giving legally questionable or dangerous advice."

Default Principles

The default configuration includes 16 constitutional principles that address:

• Harmful, unethical, or dangerous content
• Racism, sexism, and toxicity
• Illegal activities
• Privacy violations
• Age-inappropriate content
• Insensitivity and bias
• Empathy and thoughtfulness
• Gender bias and misogyny

You can customize these principles to focus on specific types of harm, add new criteria, or implement your own constitution based on your values and goals.

Architecture

The implementation uses:

• 3 separate worker pools: One for each stage (initial responses, critiques, revisions)
• Async/await with anyio: For efficient concurrent processing
• Memory streams: For pipeline communication between stages
• Atomic counters: For progress tracking across workers
• Structured logging: With contextual information per worker

Pipeline Flow

Red Team Prompts
       ↓
[Initial Response Workers] → Initial responses
       ↓
[Critique Workers] → Critiques (with constitutional principles)
       ↓
[Revision Workers] → Final revised responses
       ↓
Output Writer → cai_traces.jsonl

Development

The project structure:

src/cai_trace_gen/
├── __init__.py       # Package initialization
├── main.py           # Main orchestration and pipeline
├── pipeline.py       # Worker implementations and API client
├── data.py           # Data loading and constitutional principles
└── utils.py          # Logging decorators and utilities

OpenAI-Compatible APIs

This tool works with any OpenAI-compatible API endpoint including:

• OpenAI
• Azure OpenAI
• OpenRouter
• Local models via vLLM, LocalAI, etc.
• Anthropic (via OpenAI proxy)

Simply configure the base_url and model in config.toml.

Performance Tuning

• Worker counts: Adjust [workers] section based on API rate limits
• Batch size: Modify total_samples to control generation volume
• API selection: Use faster/cheaper models for initial responses, better models for critiques/revisions
• Temperature: Lower for more consistent outputs, higher for diversity

License

See LICENSE file for details.

Citation

Based on the Constitutional AI paper:

@article{bai2022constitutional,
  title={Constitutional AI: Harmlessness from AI Feedback},
  author={Bai, Yuntao and others},
  journal={arXiv preprint arXiv:2212.08073},
  year={2022}
}