omnigen-usf 0.1.18
Production-ready synthetic data generation framework with pipeline-based architecture, multi-provider support, auto-resume, and enterprise-grade reliability
License: MIT
Requires Python: >=3.8
Keywords: synthetic-data, data-generation, conversational-ai, text-enhancement, llm, pipeline, conversation-extension, pre-training, continual-pre-training, auto-resume, graceful-shutdown, machine-learning, ai

# OmniGen 🚀

**Enterprise-Grade Synthetic Data Generation Framework**

Production-ready with complete validation rules, auto-resume, and clean configuration.

Built by [Ultrasafe AI](https://us.inc) for production environments.

[![Version](https://img.shields.io/badge/version-0.1.18-blue.svg)](https://pypi.org/project/omnigen-usf/)
[![Python](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)

---

## What is OmniGen?

**OmniGen** is an enterprise-grade framework for generating synthetic datasets at scale—from scratch or from base data. Generate **trillions of tokens** and **billions of samples** with production-ready configuration:

### 🎯 Available Pipelines
- 💬 **Conversational Data** - Single-turn to multi-turn dialogues with tool calling support
- ✨ **Text Enhancement** - Large-scale corpus generation and text quality improvement  

### 🎓 Production Use Cases
- **Pre-Training** - Generate trillions of tokens for large language models
- **Continual Pre-Training** - Expand and enhance existing training corpora
- **Fine-Tuning** - Create instruction-following and task-specific datasets
- **Supervised Fine-Tuning (SFT)** - High-quality labeled conversational data
- **Domain Adaptation** - Generate domain-specific datasets at scale
- **Data Augmentation** - Expand existing datasets with variations

### 🏗️ Why OmniGen for Large-Scale Data Generation?
- ✅ **Scale to Billions** - Process trillions of tokens with unlimited parallelism (100, 1000, 10000+ workers)
- ✅ **Zero Duplicates** - ID-based tracking with content hashing ensures no duplicate data
- ✅ **Auto-Resume** - Seamless resume from any interruption, never lose progress
- ✅ **Cost Control** - Real-time token tracking and cost monitoring
- ✅ **Quality Assurance** - Two-level validation (message + conversation) ensures data quality
- ✅ **Rate Limiting** - Built-in concurrent call limits prevent API throttling
- ✅ **Production-Ready** - Clean configs, 80+ options, all documented
- ✅ **Fault Tolerant** - Checkpoint versioning, graceful shutdown in ~7 seconds
- ✅ **Security** - Environment variables, workspace isolation

---

## 🆕 What's New in v0.1.18

### � Universal Anthropic-Compatible Provider Support
- ✅ **Any Anthropic-compatible provider** - Use ANY provider that supports Anthropic-style SDK/API
- ✅ **User-controlled** - No auto-detection, you explicitly declare compatibility
- ✅ **Future-proof** - Works with new providers without code changes
- ✅ **Built-in support** for:
  - Official Anthropic (Claude 4)
  - MiniMax M2 via Anthropic endpoint
  - Moonshot Kimi K2 via Anthropic endpoint
  - OpenRouter (both native and Anthropic-compatible modes)

**Example - MiniMax M2:**
```yaml
providers:
  assistant_response:
    name: anthropic
    api_key: ${MINIMAX_API_KEY}
    base_url: https://api.minimax.io/anthropic
    model: MiniMax-M2
    anthropic_compatible: true  # User declares compatibility
    interleaved_thinking: true
```

### 🧠 Enhanced Reasoning Output Preservation
- ✅ **New default** - `save_all_reasoning: true` (maximum transparency)
- ✅ **Two powerful options** - Full transparency OR strategic saving
- ✅ **Flexible control** - Save reasoning at key decision points only

**Option 1 - Maximum Transparency (Default):**
```yaml
reasoning_output_rules:
  save_all_reasoning: true  # Reasoning in ALL messages
```

**Option 2 - Strategic Decision Points:**
```yaml
reasoning_output_rules:
  save_all_reasoning: false
  keep_with_tool_calls: true       # Save when calling tools
  keep_immediate_after_tool: true  # Save when processing results
  keep_last_message: true          # Save final conclusion
```

### 🐛 Bug Fixes
- ✅ **None message filtering** - Fixed edge case in interleaved thinking
- ✅ **100% test coverage** - All 31 tests passing

---

## 🚀 Available Pipelines

### 1. **conversation_extension** - Multi-Turn Dialogue Generation

Turn single-turn questions into rich multi-turn conversations with full control over turns, providers, and validation.

**Features**: Tool calls support, Auto-resume, Token tracking, Flexible turn control, Multi-provider support

### 2. **text_enhancement** - Corpus Generation & Text Quality Improvement ✨

Transform and enhance text at scale for corpus generation, quality improvement, and dataset expansion.

**Features**: Auto-resume, Graceful shutdown, Smart validation, Streaming, Token tracking, Retry logic

---

## 🚀 Key Features

### 1. **Auto-Resume Without Duplicates** 🔄
Both pipelines guarantee zero duplicates on resume:
- **ID-based tracking** with content hashing (each item has unique ID)
- **Content verification** - detects if data changed
- **Zero duplicates** - Mathematical guarantee via ID deduplication
- **Automatic skip** of already processed IDs
- **O(1) lookup** performance even for millions of items

```bash
# Just restart - auto-resumes!
$ python run_conversation_extension.py --config config.yaml
🔄 Resuming from checkpoint... ✓ Skipping 500 already processed IDs
```

### 2. **Graceful Shutdown with Real-Time Status** 🛑
Press Ctrl+C and see exactly what's happening:
- 3-step progress (Cancel → Checkpoint → Wait)
- Real-time countdown timer
- ~7 seconds maximum
- Zero data loss

```
[1/3] Canceling 1950 pending tasks... ✓ 0.08s
[2/3] Saving checkpoint... ✓ 1.42s
[3/3] Waiting for in-flight... 48 done, 3.2s left
✓ Total shutdown time: 6.82s
```

### 3. **Complete Validation System** ✅

**Message Validation** (Per-message as generated):
```yaml
message_validation:
  enabled: true
  user_checks:
    empty_content: true          # No empty messages
    min_length: false            # Optional min length
    no_duplicate_in_history: true # No repetition
    regex_match: false           # Optional pattern matching
    contains: false              # Optional required strings
```

**Quality Validation** (Entire conversation after generation):
```yaml
quality_validation:
  enabled: true
  checks:
    empty_content: true
    repeated_messages: true
    alternation: true            # Proper user/assistant order
    tool_calls: true             # Tool call consistency
```

### 4. **Rate Limiting & Performance** ⚡

```yaml
providers:
  user_followup:
    max_concurrent_calls: 70     # Prevent API throttling

generation:
  parallel_workers: 70           # Match concurrent calls
  track_tokens: true             # Monitor usage
  token_pricing:
    input_cost_per_million: 1.0  # Track costs
```

### 5. **Production-Ready Config** 📋

Clean, focused configuration with inline documentation:
- ✅ 241 lines (down from 360+)
- ✅ Essential features enabled by default
- ✅ Optional features commented out
- ✅ Based on real production usage
- ✅ Every option explained

### 6. **Tool Calls Support** 🛠️
Full OpenAI/Anthropic API compliance - Assistant messages can have empty content if `tool_calls` present.

### 7. **Real-Time Token Tracking** 💰
- Captures actual token usage from API responses
- Tracks per generation call
- Aggregates per conversation/dataset
- Configurable cost monitoring
- Saves detailed token data in output files


---

## Quick Start

### 1. Install

```bash
pip install omnigen-usf
```

### 2. Get Production-Ready Config

```bash
# Conversation Extension
# Option A: Full config with all options documented (242 lines)
cp examples/conversation_extension/config.yaml my_config.yaml

# Option B: Minimal config with only essentials (132 lines)
cp examples/conversation_extension/config_minimal.yaml my_config.yaml

# Text Enhancement
# Option A: Full config with all options documented (374 lines)
cp examples/text_enhancement/config.yaml my_config.yaml

# Option B: Minimal config with only essentials (133 lines)
cp examples/text_enhancement/config_minimal.yaml my_config.yaml
```

### 3. Customize & Run

**Update config** with your settings:
```yaml
workspace_id: "my-project-v1"

providers:
  user_followup:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}      # Use environment variable
    model: usf-mini
    max_concurrent_calls: 70            # Your rate limit

generation:
  num_conversations: 0              # Process all data
  parallel_workers: 70              # Match rate limit
  track_tokens: true

base_data:
  file_path: data/input.jsonl

storage:
  output_file: data/output.jsonl

checkpoint:
  enabled: true
  checkpoint_file: data/checkpoint.json
```

**Set environment variable**:
```bash
export USF_API_KEY="your-api-key"
```

**Run**:
```bash
python run_conversation_extension.py --config my_config.yaml
```

### 4. Alternative: Programmatic Usage

#### Option A: Using YAML Configuration File

Create a `config.yaml` file:

```yaml
# NEW: Minimal config with smart defaults!
providers:
  user_followup:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}
    # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓
  assistant_response:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}
    # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓

generation:
  num_conversations: 100
  turn_range: {min: 3, max: 8}

base_data:
  source_type: file
  file_path: base_data.jsonl

storage:
  type: jsonl
  output_file: output.jsonl
```

Then load and run:

```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfig,
    ConversationExtensionPipeline
)

# Load configuration from YAML file
config = ConversationExtensionConfig.from_yaml('config.yaml')

# Run the pipeline
pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

#### Option B: Using Programmatic Configuration (Python)

```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfigBuilder,
    ConversationExtensionPipeline
)

# Configure the pipeline programmatically
config = (ConversationExtensionConfigBuilder()
    # User followup generator - minimal config!
    .add_provider(
        role='user_followup',
        name='ultrasafe',
        api_key='your-api-key'
        # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓
    )
    # Assistant response generator
    .add_provider(
        role='assistant_response',
        name='ultrasafe',
        api_key='your-api-key'
        # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓
    )
    # Generation settings
    .set_generation(
        num_conversations=100,
        turn_range=(3, 8)  # 3-8 turns per conversation
    )
    # Input data
    .set_data_source(
        source_type='file',
        file_path='base_data.jsonl'
    )
    # Output
    .set_storage(
        type='jsonl',
        output_file='output.jsonl'
    )
    .build()
)

# Run the pipeline
pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

### 4. Get Results

Your generated conversations will be in `output.jsonl`:

```jsonl
{
  "id": 0,
  "conversations": [
    {"role": "user", "content": "How do I learn Python?"},
    {"role": "assistant", "content": "Great choice! Start with the basics..."},
    {"role": "user", "content": "What resources do you recommend?"},
    {"role": "assistant", "content": "I recommend these resources..."},
    {"role": "user", "content": "How long will it take?"},
    {"role": "assistant", "content": "With consistent practice..."}
  ],
  "num_turns": 3,
  "success": true
}
```

---

## Supported AI Providers

**NEW: Smart defaults automatically applied!** Only specify `name` and `api_key` - model, temperature, and max_tokens use optimized defaults.

| Provider | Default Model | Default Temp | Default Tokens |
|----------|---------------|--------------|----------------|
| **Ultrasafe AI** | `usf-mini` | `0.7` | `4096` |
| **OpenAI** | `gpt-4-turbo` | `0.7` | `4096` |
| **Anthropic** | `claude-3-5-sonnet-20241022` | `0.7` | `4096` |
| **OpenRouter** | `openai/gpt-4-turbo` | `0.7` | `4096` |

**Override defaults as needed:**
```yaml
providers:
  user_followup:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}
    model: usf-mini
    temperature: 0.9  # Override only what you need
```

### Mix Different Providers

```python
config = (ConversationExtensionConfigBuilder()
    .add_provider('user_followup', 'ultrasafe', api_key, 'usf-mini')
    .add_provider('assistant_response', 'ultrasafe', api_key, 'usf-mini')
    # ... rest of config
    .build()
)
```

---

## Advanced Features

### 🔄 Checkpoint & Resume System

**NEW!** Automatic checkpoint/resume functionality prevents data loss and duplication from any interruption.

#### Overview

The checkpoint system automatically saves progress and enables seamless resume from where you left off - handling manual stops (Ctrl+C), errors, rate limits, server failures, and crashes without losing data or creating duplicates.

#### How It Works

**Conversation Tracking:**
- **Position-based ID**: Index in base data file (for ordering)
- **Content Hash**: SHA256 hash of conversation (for deduplication)
- **Combined Key**: `{position}_{hash[:8]}` for unique identification

**Checkpoint Saves:**
- Automatically after every N conversations (configurable)
- On graceful shutdown (Ctrl+C, SIGTERM)
- Uses atomic file operations (write temp → sync → atomic rename)
- Validates input file integrity with SHA256 hash

#### Configuration

```yaml
checkpoint:
  enabled: true                    # Enable checkpoint/resume
  checkpoint_file: "checkpoint.json"  # Path to checkpoint file
  batch_save_items: 100            # Save every 100 items
  batch_save_seconds: 10           # OR every 10 seconds (whichever first)
  validate_input_hash: true        # Verify input hasn't changed
  resume_mode: "auto"              # auto | manual | fresh
```

**Configuration Options:**

| Option | Description | Default |
|--------|-------------|---------|
| `enabled` | Enable checkpoint/resume | `true` |
| `checkpoint_file` | Path to checkpoint file | `workspaces/{workspace_id}/checkpoint.json` |
| `batch_save_items` | Save checkpoint every N items | `100` |
| `batch_save_seconds` | Save checkpoint every N seconds (OR condition) | `10` |
| `validate_input_hash` | Verify input file unchanged on resume | `true` |
| `resume_mode` | Resume behavior: `auto`, `manual`, or `fresh` | `auto` |

#### Usage Examples

**YAML Configuration:**
```yaml
workspace_id: "production_run"

checkpoint:
  enabled: true
  checkpoint_file: "checkpoint.json"
  batch_save_items: 100
  batch_save_seconds: 10
  validate_input_hash: true
  resume_mode: "auto"

providers:
  user_followup:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}
  assistant_response:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}

generation:
  num_conversations: 1000
  turn_range: {min: 3, max: 8}
  parallel_workers: 10

base_data:
  source_type: file
  file_path: input.jsonl

storage:
  type: jsonl
  output_file: output.jsonl
```

**Programmatic Configuration:**
```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfigBuilder,
    ConversationExtensionPipeline
)

config = (ConversationExtensionConfigBuilder()
    .add_provider('user_followup', 'ultrasafe', api_key, 'usf-mini')
    .add_provider('assistant_response', 'ultrasafe', api_key, 'usf-mini')
    .set_generation(num_conversations=1000, turn_range=(3, 8))
    .set_data_source('file', file_path='input.jsonl')
    .set_storage('jsonl', output_file='output.jsonl')
    .set_checkpoint(
        enabled=True,
        checkpoint_file='checkpoint.json',
        batch_save_items=100,
        batch_save_seconds=10,
        validate_input_hash=True,
        resume_mode='auto'
    )
    .build()
)

pipeline = ConversationExtensionPipeline(config)
pipeline.run()  # Can interrupt and resume anytime!
```

#### Resume Behavior

**First Run (No Checkpoint):**
```
==========================================
CONVERSATION EXTENSION PIPELINE
==========================================
Base conversations: 1000
Generating: 1000
Parallel workers: 10
==========================================
Generating: 15%|███░░░░░░░░░| 150/1000 [✓120 ⚠20 ✗10]
^C
⚠️  Shutdown signal received. Saving checkpoint...
✓ Final checkpoint saved
⚠️  Interrupted. Progress saved in checkpoint.
```

**Resume Run (Checkpoint Exists):**
```
==========================================
RESUMING FROM CHECKPOINT
==========================================
Previous Run: 2025-10-04 16:30:00 UTC
Already Processed: 150 (✓120 ⚠20 ✗10 ~0)
Resuming Partials: 20
Remaining: 850
Parallel workers: 10
==========================================
Generating: 100%|████████████| 1000/1000 [✓850 ⚠100 ✗50]

==========================================
GENERATION COMPLETE
==========================================
✓ Complete:   850
⚠ Partial:    100
✗ Failed:      50
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
💾 Saved:     950  (95.0%)
==========================================
```

#### Output Format

Both partial and complete conversations saved to **same output file** with status indicators:

```jsonl
{"id": 0, "status": "completed", "is_complete": true, "conversations": [...], "num_turns": 5}
{"id": 1, "status": "partial", "is_complete": false, "conversations": [...], "num_turns": 3}
{"id": 2, "status": "completed", "is_complete": true, "conversations": [...], "num_turns": 7}
```

**Status Values:**
- `"completed"` - Successfully generated all requested turns
- `"partial"` - Interrupted during generation, may be incomplete
- `"failed"` - Error occurred, saved to failed file

**Filtering Results:**
```bash
# Get only completed conversations
jq 'select(.status == "completed")' output.jsonl > completed.jsonl

# Get partial conversations (to review/complete)
jq 'select(.status == "partial")' output.jsonl > partial.jsonl

# Count by status
jq -s 'group_by(.status) | map({status: .[0].status, count: length})' output.jsonl
```

#### Key Features

✅ **Zero Data Loss**
- Progress saved automatically at configurable intervals
- Every conversation saved before moving to next
- Atomic file operations prevent corruption

✅ **No Duplicates**
- Hybrid identification (position + content hash)
- Skips already processed conversations automatically
- Prevents reprocessing on resume

✅ **Partial Resume**
- Continues incomplete multi-turn conversations
- Resumes from exact turn where interrupted
- Preserves all generated content

✅ **Universal Coverage**
- Handles ALL interruption types:
  - Manual stops (Ctrl+C)
  - Server/connection failures
  - Rate limit errors
  - System crashes
  - Any other errors

✅ **Input Validation**
- SHA256 hash verification of input file
- Detects if input changed between runs
- Prevents processing wrong data

✅ **Safety Features**
- Atomic file operations (corruption-proof)
- Automatic backup of corrupted checkpoints
- Graceful shutdown handlers (SIGINT/SIGTERM)

#### Checkpoint Structure

```json
{
  "version": "1.0",
  "run_id": "unique-uuid",
  "started_at": "2025-10-04T20:00:00.000Z",
  "last_checkpoint_at": "2025-10-04T20:15:30.000Z",
  "base_data": {
    "source_type": "file",
    "file_path": "input.jsonl",
    "file_hash": "sha256-hash",
    "total_available": 1000
  },
  "progress": {
    "total_processed": 150,
    "completed": 120,
    "partial": 20,
    "failed": 10,
    "last_position": 149
  },
  "processed_records": [
    {
      "position": 0,
      "content_hash": "a1b2c3d4",
      "status": "completed",
      "turns_generated": 5,
      "processed_at": "2025-10-04T20:01:00.000Z"
    }
  ],
  "partial_states": {
    "5_b2c3d4e5": {
      "position": 5,
      "conversation": [...],
      "turns_completed": 3,
      "target_turns": 5
    }
  }
}
```

#### Best Practices

1. **Set Appropriate Save Frequency**
   ```yaml
   checkpoint:
     batch_save_items: 100      # Save every 100 items
     batch_save_seconds: 10     # OR every 10 seconds
   ```
   - Lower items (50-100): Better safety, slight overhead
   - Higher items (500-1000): Better performance, less frequent saves
   - Time-based fallback ensures saves even if processing is slow

2. **Use Workspace Isolation**
   ```python
   config = ConversationExtensionConfigBuilder(
       workspace_id="unique_project_id"  # Prevents conflicts
   )
   ```

3. **Monitor Checkpoint Size**
   - Large checkpoints (>100MB) may slow saves
   - Consider processing in batches for huge datasets

4. **Backup Important Checkpoints**
   ```bash
   cp checkpoint.json checkpoint.backup.json
   ```

5. **Clean Up After Completion**
   ```bash
   # After successful run
   rm checkpoint.json
   ```

#### Troubleshooting

**Problem: Checkpoint not resuming**
- Check `checkpoint.enabled` is `true`
- Verify checkpoint file exists at specified path
- Ensure `resume_mode` is not set to `"fresh"`
- Review logs for validation errors

**Problem: Input file changed warning**
- If intentional: Set `validate_input_hash: false`
- If unintentional: Restore original input file
- Or delete checkpoint to start fresh

**Problem: Partial conversations not resuming**
- Verify `partial_states` exists in checkpoint
- Check conversation has `_position` and `_content_hash`
- Review logs for partial state loading errors

**Problem: Duplicate conversations**
- Checkpoint prevents this automatically
- Check both `position` and `content_hash` are tracked
- Verify hybrid identification is working

#### Performance Impact

- **Checkpoint overhead**: ~10-50ms per save (amortized)
- **Memory usage**: Minimal (lightweight data structure)
- **Disk space**: ~1-2KB per conversation in checkpoint
- **Resume speed**: Instant (IDs loaded into memory set)

#### Example: Complete Workflow

```python
import os
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfigBuilder,
    ConversationExtensionPipeline
)

# Build configuration
config = (ConversationExtensionConfigBuilder(workspace_id="production_v1")
    .add_provider(
        role="user_followup",
        name="ultrasafe",
        api_key=os.getenv("ULTRASAFE_API_KEY"),
        model="usf-mini"
    )
    .add_provider(
        role="assistant_response",
        name="ultrasafe",
        api_key=os.getenv("ULTRASAFE_API_KEY"),
        model="usf-mini"
    )
    .set_generation(
        num_conversations=10000,
        turn_range=(3, 8),
        parallel_workers=20
    )
    .set_data_source(
        source_type="file",
        file_path="data/base_conversations.jsonl"
    )
    .set_storage(
        type="jsonl",
        output_file="output.jsonl"
    )
    .set_checkpoint(
        enabled=True,
        checkpoint_file="checkpoint.json",
        batch_save_items=50,
        batch_save_seconds=10,
        validate_input_hash=True,
        resume_mode="auto"
    )
    .build()
)

# Run with automatic resume
pipeline = ConversationExtensionPipeline(config)
pipeline.run()

# Can be interrupted and resumed multiple times
# Progress is never lost!
```

📝 **See also**: [`examples/checkpoint_resume_example.py`](examples/conversation_extension/checkpoint_resume_example.py) for a complete working example with checkpoint inspection utilities.

---

### Multi-Tenant SaaS Support

Perfect for platforms serving multiple users concurrently:

```python
# Each user gets isolated workspace
workspace_id = f"user_{user_id}_session_{session_id}"

config = (ConversationExtensionConfigBuilder(workspace_id=workspace_id)
    .add_provider('user_followup', 'ultrasafe', shared_api_key, 'usf-mini')
    .add_provider('assistant_response', 'ultrasafe', shared_api_key, 'usf-mini')
    .set_storage('jsonl', output_file='output.jsonl')  # Auto-isolated
    .build()
)

# Storage automatically goes to: workspaces/{workspace_id}/output.jsonl
```

### Parallel Dataset Generation

```python
from concurrent.futures import ProcessPoolExecutor

def process_dataset(input_file, output_file):
    config = (ConversationExtensionConfigBuilder()
        .add_provider('user_followup', 'ultrasafe', api_key, 'usf-mini')
        .add_provider('assistant_response', 'ultrasafe', api_key, 'usf-mini')
        .set_data_source('file', file_path=input_file)
        .set_storage('jsonl', output_file=output_file)
        .build()
    )
    ConversationExtensionPipeline(config).run()

# Process 3 datasets in parallel
with ProcessPoolExecutor(max_workers=3) as executor:
    executor.submit(process_dataset, 'data1.jsonl', 'out1.jsonl')
    executor.submit(process_dataset, 'data2.jsonl', 'out2.jsonl')
    executor.submit(process_dataset, 'data3.jsonl', 'out3.jsonl')
```

## 📝 Configuration Methods

OmniGen supports **multiple ways** to configure your pipeline. Choose the method that best fits your workflow:

### Method 1: YAML Configuration File ⭐ (Recommended)

The most flexible and maintainable approach:

**Step 1:** Create `config.yaml`
```yaml
# NEW: Minimal configuration with smart defaults
providers:
  user_followup:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}  # Only name and API key required!
    # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓
  assistant_response:
    name: ultrasafe
    api_key: ${ULTRASAFE_API_KEY}
    # Defaults: usf-mini, 0.7 temp, 4096 tokens ✓

generation:
  num_conversations: 100
  turn_range: {min: 3, max: 8}

base_data:
  source_type: file
  file_path: base_data.jsonl

storage:
  type: jsonl
  output_file: output.jsonl
```

**Step 2:** Load and run
```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfig,
    ConversationExtensionPipeline
)

# Load from YAML file
config = ConversationExtensionConfig.from_yaml('config.yaml')

# Run pipeline
pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

**✅ Benefits:**
- Easy to version control and share
- Environment variable support with `${VAR_NAME}`
- No code changes needed for config updates
- Clear, readable configuration

### Method 2: Programmatic Configuration (Python)

Build configuration directly in code:

```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfigBuilder,
    ConversationExtensionPipeline
)

config = (ConversationExtensionConfigBuilder()
    .add_provider('user_followup', 'ultrasafe', 'api-key', 'usf-mini')
    .add_provider('assistant_response', 'ultrasafe', 'api-key', 'usf-mini')
    .set_generation(num_conversations=100, turn_range=(3, 8))
    .set_data_source('file', file_path='base_data.jsonl')
    .set_storage('jsonl', output_file='output.jsonl')
    .build()
)

pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

**✅ Benefits:**
- Type safety and IDE autocomplete
- Dynamic configuration based on runtime
- No external files needed
- Easy application integration

### Method 3: Dictionary Configuration

Create from Python dictionary:

```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfig,
    ConversationExtensionPipeline
)

config_dict = {
    'providers': {
        'user_followup': {'name': 'ultrasafe', 'api_key': 'key', 'model': 'usf-mini'},
        'assistant_response': {'name': 'ultrasafe', 'api_key': 'key', 'model': 'usf-mini'}
    },
    'generation': {'num_conversations': 100, 'turn_range': {'min': 3, 'max': 8}},
    'base_data': {'source_type': 'file', 'file_path': 'base_data.jsonl'},
    'storage': {'type': 'jsonl', 'output_file': 'output.jsonl'}
}

config = ConversationExtensionConfig.from_dict(config_dict)
pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

**✅ Benefits:**
- Load from JSON files or APIs
- Easy programmatic modification
- Flexible for dynamic scenarios

### Method 4: Hybrid Approach

Combine YAML with programmatic overrides:

```python
from omnigen.pipelines.conversation_extension import (
    ConversationExtensionConfig,
    ConversationExtensionPipeline
)

# Load base config from YAML
config = ConversationExtensionConfig.from_yaml('base_config.yaml')

# Modify specific settings
config_dict = config.to_dict()
config_dict['generation']['num_conversations'] = 500  # Override
config_dict['storage']['output_file'] = f'output_{user_id}.jsonl'  # Dynamic

# Rebuild and run
config = ConversationExtensionConfig.from_dict(config_dict)
pipeline = ConversationExtensionPipeline(config)
pipeline.run()
```

**✅ Benefits:**
- Base configuration in YAML
- Runtime customization in code
- Best of both worlds

### Environment Variables in YAML

Use `${VARIABLE_NAME}` syntax:

```yaml
providers:
  user_followup:
    name: ${PROVIDER_NAME}           # From environment
    api_key: ${ULTRASAFE_API_KEY}    # From environment
    model: ${USER_MODEL}             # From environment

base_data:
  file_path: ${INPUT_PATH}           # From environment

storage:
  output_file: ${OUTPUT_PATH}        # From environment
```

Set variables:
```bash
export PROVIDER_NAME="ultrasafe"
export ULTRASAFE_API_KEY="your-key"
export USER_MODEL="usf-mini"
export INPUT_PATH="data/input.jsonl"
export OUTPUT_PATH="data/output.jsonl"

# Then run
python your_script.py
```

### CLI Usage with Config File

Run directly from command line:

```bash
# Using YAML config
omnigen conversation-extension --config config.yaml

# With overrides
omnigen conversation-extension \
  --config config.yaml \
  --num-conversations 500 \
  --output custom_output.jsonl
```

---

---

## 📖 Complete Configuration Reference

### All Configuration Options Explained

Below is a comprehensive YAML configuration showing **ALL** available options with detailed explanations:

```yaml
# ==============================================================================
# WORKSPACE ISOLATION (Optional)
# ==============================================================================
# Unique ID for multi-tenant environments - auto-isolates all output files
workspace_id: "user_123_session_abc"

# ==============================================================================
# PROVIDERS - AI Model Configuration
# ==============================================================================
# Configure different AI providers for each role
# Each role can use a different provider/model combination

providers:
  # Provider for generating user follow-up questions
  user_followup:
    name: ultrasafe
    api_key: ${API_KEY}          # Required: Use env var ${VAR_NAME} or direct key
    
    # Optional - Smart defaults applied if not specified:
    model: usf-mini              # Default varies by provider
    temperature: 0.7             # Default: 0.7
    max_tokens: 2048             # Default: 4096 (overridden here)
    timeout: 300                 # Default: 300
    max_retries: 5               # Default: 5
    retry_delay: 2               # Default: 2
  
  # Provider for generating assistant responses
  assistant_response:
    name: ultrasafe              # Can use different provider than user_followup
    api_key: ${API_KEY}          # Only name and api_key are required!
    max_tokens: 8192             # Override default (4096) for detailed responses
    # model, temperature, timeout, retries use defaults ✓

# PROVIDER OPTIONS:
# ----------------
# ultrasafe:
#   models: usf-mini, usf-max
#
# openai:
#   models: gpt-4-turbo, gpt-4, gpt-3.5-turbo, gpt-4o, gpt-4o-mini
#
# anthropic:
#   models: claude-3-5-sonnet-20241022, claude-3-opus-20240229,
#           claude-3-sonnet-20240229, claude-3-haiku-20240307
#
# openrouter:
#   models: Any OpenRouter supported model
# ==============================================================================
# TOKEN TRACKING (NEW!)
# ==============================================================================
generation:
  track_tokens: true                   # Enable real-time token tracking from API
  token_pricing:                       # Optional: for cost display in console
    input_cost_per_million: 0.35       # Your input token pricing
    output_cost_per_million: 2.50      # Your output token pricing

# Token tracking captures ACTUAL usage from API responses:
# - Per generation call (user_followup, assistant_response)
# - Per conversation (aggregated)
# - Total dataset (summed across all conversations)
#
# Output includes detailed token breakdown:
# {
#   "tokens": {
#     "input_tokens": 24129,
#     "output_tokens": 5924,
#     "total_tokens": 30053,
#     "per_generation": [
#       {"type": "user_followup", "input_tokens": 3150, "output_tokens": 41},
#       {"type": "assistant_response", "input_tokens": 2632, "output_tokens": 927}
#     ]
#   }
# }
#
# NOTE: Token data is saved to output files
# Cost calculation is optional and only shown in console

#   base_url: https://openrouter.ai/api/v1 (optional)

# ==============================================================================
# GENERATION SETTINGS
# ==============================================================================
generation:
  num_conversations: 100           # Total conversations to generate
                                   # Use 0 or omit to process ALL available conversations
  
  turn_range:                      # Number of turns per conversation
    min: 3                         # Minimum turns
    max: 8                         # Maximum turns
  
  parallel_workers: 10             # Concurrent workers (balance speed vs rate limits)
  
  # Extension behavior for multi-turn input
  extension_mode: "smart"          # Options: "smart" | "legacy"
  # - smart: Intelligently handle multi-turn conversations
  # - legacy: Always extract first user message only
  
  skip_invalid: true               # Skip invalid patterns (recommended: true)
  
  # Turn calculation method
  turn_calculation: "additional"   # Options: "additional" | "total"
  # - additional: Add NEW turns on top of existing (default)
  # - total: Keep total turns within range (never removes existing)

# ==============================================================================
# DATA SOURCE CONFIGURATION
# ==============================================================================
base_data:
  enabled: true                    # Enable base data loading
  
  # OPTION 1: Local File
  source_type: file                # Use local JSONL/JSON file
  file_path: data/input.jsonl      # Path to file
  format: conversations            # JSON key containing conversation array
  shuffle: false                   # Shuffle data before processing
  
  # OPTION 2: HuggingFace Dataset
  # source_type: huggingface       # Use HuggingFace dataset
  # hf_dataset: username/dataset   # HuggingFace dataset path
  # hf_split: train                # Dataset split: train, test, validation
  # hf_token: ${HF_TOKEN}          # HuggingFace API token (if private)
  # hf_streaming: false            # Stream dataset (for large datasets)
  # format: conversations          # Field name in dataset
  # shuffle: true                  # Shuffle after loading

# ==============================================================================
# STORAGE CONFIGURATION
# ==============================================================================
storage:
  type: jsonl                      # Options: jsonl | mongodb
  
  # JSONL Storage (Default)
  output_file: output.jsonl        # Successful conversations
  partial_file: partial.jsonl      # Partial/incomplete conversations (legacy)
  failed_file: failed.jsonl        # Failed conversations
  
  # MongoDB Storage (Alternative)
  # type: mongodb
  # mongodb:
  #   connection_string: mongodb://localhost:27017
  #   database: omnigen
  #   collection: conversations
  #   output_collection: output          # Successful
  #   partial_collection: partial        # Partial
  #   failed_collection: failed          # Failed

# ==============================================================================
# CHECKPOINT/RESUME CONFIGURATION (NEW!)
# ==============================================================================
checkpoint:
  enabled: true                    # Enable automatic checkpoint/resume
  checkpoint_file: "workspaces/{workspace_id}/checkpoint.json"
  batch_save_items: 100            # Save checkpoint every N items
  batch_save_seconds: 10           # OR every N seconds (whichever first)
  validate_input_hash: true        # Verify input file hasn't changed on resume
  resume_mode: "auto"              # Options: "auto" | "manual" | "fresh"
  
  # Features:
  # - Zero data loss from interruptions
  # - No duplicate processing
  # - Partial conversation resume
  # - Handles: Ctrl+C, errors, rate limits, crashes

# ==============================================================================
# DATETIME CONFIGURATION (Optional)
# ==============================================================================
datetime_config:
  enabled: true                    # Enable datetime generation
  mode: random_from_range          # Options: random_from_range | current | fixed
  timezone: UTC                    # Timezone (UTC, America/New_York, Asia/Dubai, etc.)
  format: "%Y-%m-%d %H:%M:%S"      # Python strftime format
  
  # For random_from_range mode
  range:
    start: "2024-01-01 00:00:00"   # Start datetime
    end: "2024-12-31 23:59:59"     # End datetime
  
  # For fixed mode
  # fixed_datetime: "2024-06-15 12:00:00"

# ==============================================================================
# SYSTEM MESSAGES (Optional)
# ==============================================================================
system_messages:
  # Prepend system message to every conversation
  prepend_always:
    enabled: true
    content: "You are a helpful AI assistant. Current time: {current_datetime} ({timezone})."
  
  # Append system message to every conversation
  append_always:
    enabled: false
    content: "Remember to be concise and helpful."
  
  # Add system message only if none exists
  add_if_missing:
    enabled: false
    content: "You are an AI assistant."

# Available variables in system messages:
# - {current_datetime}: Generated datetime
# - {timezone}: Configured timezone
# - {workspace_id}: Current workspace ID

# ==============================================================================
# GENERATION-ONLY SYSTEM MESSAGES (Optional - NEW FEATURE!)
# ==============================================================================
# These system messages are used during generation but NOT saved to the dataset.
# Perfect for providing internal guidance without polluting training data.

generation_system_messages:
  assistant_response:
    enabled: true
    content: |
      INTERNAL GENERATION GUIDELINES (not saved to dataset):
      
      Quality Standards:
      - Accuracy: Provide factually correct information
      - Clarity: Use simple, clear language
      - Completeness: Address all aspects of the question
      - Examples: Include concrete examples when helpful
      
      Formatting:
      - Use markdown for better readability
      - Break long responses into sections
      - Use bullet points for lists
      
      Tone:
      - Professional yet friendly
      - Encouraging and supportive
      - Clear and direct

# Available variables (same as system_messages):
# - {current_datetime}: Generated datetime
# - {timezone}: Configured timezone

# ==============================================================================
# CUSTOM PROMPTS (Optional - Smart Defaults Provided!)
# ==============================================================================
# NOTE: Prompts are completely OPTIONAL. The system automatically uses optimized
# default prompts if you don't specify any. Only add this section if you want
# to customize the prompt behavior.

prompts:
  # Custom prompt for user follow-up generation
  followup_question: |
    ## Your Task
    Generate an intelligent follow-up user question based on conversation history.
    
    ### CONVERSATION HISTORY:
    {history}
    
    ### INSTRUCTIONS:
    - Generate a meaningful follow-up question
    - Be conversational