codeoptix 0.1.0__py3-none-any.whl

codeoptix/vendor/bloom/scripts/step2_ideation.py

@@ -0,0 +1,667 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import re
+import sys
+from pathlib import Path
+
+# CodeOptiX: Fixed imports for vendored Bloom
+from codeoptix.vendor.bloom.globals import models
+from codeoptix.vendor.bloom.utils import (get_reasoning_effort,
+                   parse_message, get_results_dir,
+                   litellm_chat, load_config, load_understanding_results,
+                   save_results_locally, setup_project_path)
+from codeoptix.vendor.bloom.prompts.step2_ideation import (make_system_prompt,
+                                     make_all_scenarios_prompt,
+                                     make_variation_system_prompt, make_variation_prompt)
+
+# Simple debug print function (replaces bloom.debug_print)
+def debug_print(*args, **kwargs):
+    """Debug print function for Bloom scripts."""
+    import logging
+    logging.debug(*args, **kwargs)
+
+def get_model_max_output_tokens(model_id):
+    """
+    Get the maximum output tokens for a model using litellm.
+
+    Note: Some older models (Claude 3.7) report inflated limits. We cap those based on observed behavior.
+    For newer models, we trust litellm's reported limits.
+    """
+    try:
+        import litellm
+        model_info = litellm.get_model_info(model_id)
+        max_output = model_info.get("max_output_tokens")
+
+        if max_output and isinstance(max_output, int) and max_output > 0:
+            debug_print(f"   litellm reports max_output_tokens: {max_output}")
+
+            # Only apply caps for models with known discrepancies
+            model_lower = model_id.lower()
+
+            # Claude 3.7 and earlier 3.x models report inflated limits - cap at 8K (observed limit ~6-8K)
+            if ("claude-3-7" in model_lower or "claude-3-5" in model_lower or
+                "claude-3.7" in model_lower or "claude-3.5" in model_lower or
+                "claude-3-opus" in model_lower or "claude-3-sonnet" in model_lower):
+                practical_max = min(max_output, 8192)
+                if practical_max < max_output:
+                    debug_print(f"   Applied cap for Claude 3.x: {practical_max} (observed limit)")
+            else:
+                # For all other models (Claude 4+, GPT, Gemini, etc.), trust litellm's reported limit
+                practical_max = max_output
+
+            return practical_max
+
+    except Exception as e:
+        debug_print(f"⚠️  Could not get max output tokens from litellm for {model_id}: {e}")
+
+    # Fallback to conservative defaults if litellm lookup fails
+    debug_print(f"⚠️  Using fallback max output tokens for {model_id}")
+    return 8192  # Conservative default
+
+
+def calculate_batch_size(total_scenarios, model_id, modality="conversation", reasoning_effort="none", safety_margin=0.8):
+    """
+    Calculate optimal batch size for scenario generation.
+
+    Args:
+        total_scenarios: Total number of scenarios to generate
+        model_id: Model identifier
+        modality: "conversation" or "simenv"
+        reasoning_effort: Extended thinking level ("none", "low", "medium", "high")
+        safety_margin: Safety factor to avoid hitting limits (0.8 = use 80% of max)
+
+    Returns:
+        batch_size: Number of scenarios per batch
+        num_batches: Number of batches needed
+    """
+    # Get model's max output tokens
+    max_output_tokens = get_model_max_output_tokens(model_id)
+
+    # Subtract thinking tokens based on reasoning effort
+    # Extended thinking uses FIXED token budgets (not proportional):
+    # Source: https://docs.litellm.ai/docs/providers/anthropic
+    thinking_tokens_fixed = {
+        "none": 0,
+        "low": 1024,    # Fixed 1,024 token budget
+        "medium": 2048, # Fixed 2,048 token budget (updated from docs)
+        "high": 4096    # Fixed 4,096 token budget (updated from docs)
+    }.get(reasoning_effort, 0)
+
+    # Content tokens = total - fixed thinking budget
+    content_tokens = max_output_tokens - thinking_tokens_fixed
+
+    debug_print(f"   Thinking allocation ({reasoning_effort}): {thinking_tokens_fixed} tokens (fixed)")
+    debug_print(f"   Content tokens available: {content_tokens}")
+
+    # Estimate tokens per scenario (based on observed actual usage)
+    # SimEnv scenarios: ~850 tokens on average (observed from actual runs)
+    # Conversation scenarios: ~500 tokens on average estimate
+    # Adding 20% buffer for longer scenarios
+    tokens_per_scenario = 1000 if modality == "simenv" else 600
+
+    # Reserve buffer for formatting overhead and longer-than-expected scenarios
+    overhead_tokens = 2000
+
+    # Calculate how many scenarios fit in one batch
+    usable_tokens = int((content_tokens - overhead_tokens) * safety_margin)
+    batch_size = max(1, usable_tokens // tokens_per_scenario)
+
+    # Don't make batches larger than total scenarios
+    batch_size = min(batch_size, total_scenarios)
+
+    # Calculate number of batches
+    num_batches = (total_scenarios + batch_size - 1) // batch_size  # Ceiling division
+
+    debug_print(f"📊 Batch calculation:")
+    debug_print(f"   Model max output tokens: {max_output_tokens}")
+    debug_print(f"   Tokens per scenario ({modality}): {tokens_per_scenario}")
+    debug_print(f"   Usable tokens (with {int(safety_margin*100)}% safety margin): {usable_tokens}")
+    debug_print(f"   Calculated batch size: {batch_size} scenarios")
+    debug_print(f"   Number of batches: {num_batches}")
+
+    return batch_size, num_batches
+
+
+def parse_scenarios_response(response_text):
+    """Parse the scenarios response to extract individual scenarios with tools and thinking."""
+    scenarios = []
+    overall_thinking = ""
+    
+    # Extract overall thinking from <thinking> tags
+    thinking_match = re.search(r'<thinking>(.*?)</thinking>', response_text, re.DOTALL)
+    if thinking_match:
+        overall_thinking = thinking_match.group(1).strip()
+    
+    # First try to extract scenarios from <scenario> tags
+    scenario_matches = re.findall(r'<scenario>(.*?)</scenario>', response_text, re.DOTALL)
+    
+    if scenario_matches:
+        # If we found scenario tags, use those
+        for match in scenario_matches:
+            scenario_text = match.strip()
+            
+            # Extract tool signatures from this scenario
+            tool_signatures = re.findall(r'<tool_signature>(.*?)</tool_signature>', scenario_text, re.DOTALL)
+            
+            # Remove tool signatures from scenario description to get clean text
+            clean_scenario = re.sub(r'<tool_signature>.*?</tool_signature>', '', scenario_text, flags=re.DOTALL)
+            clean_scenario = re.sub(r'\s+', ' ', clean_scenario.strip())
+            
+            if clean_scenario:
+                scenario_data = {
+                    "description": clean_scenario,
+                    "tools": tool_signatures,
+                    "thinking": overall_thinking
+                }
+                scenarios.append(scenario_data)
+    else:
+        # If no scenario tags, try to split by "Scenario X:" patterns
+        scenario_splits = re.split(r'\*\*Scenario \d+:', response_text)
+        if len(scenario_splits) > 1:
+            # Skip the first split (it's before the first scenario)
+            for split in scenario_splits[1:]:
+                scenario_text = split.strip()
+                
+                # Extract tool signatures from this scenario
+                tool_signatures = re.findall(r'<tool_signature>.*?</tool_signature>', scenario_text, re.DOTALL)
+                
+                # Remove tool signatures from scenario description
+                clean_scenario = re.sub(r'<tool_signature>.*?</tool_signature>', '', scenario_text, flags=re.DOTALL)
+                clean_scenario = re.sub(r'\s+', ' ', clean_scenario.strip())
+                
+                if clean_scenario:
+                    scenario_data = {
+                        "description": clean_scenario,
+                        "tools": tool_signatures,
+                        "thinking": overall_thinking
+                    }
+                    scenarios.append(scenario_data)
+        else:
+            # Fallback: try to split by numbered patterns
+            scenario_splits = re.split(r'\d+\.\s*', response_text)
+            if len(scenario_splits) > 1:
+                for split in scenario_splits[1:]:
+                    scenario_text = split.strip()
+                    
+                    # Extract tool signatures from this scenario
+                    tool_signatures = re.findall(r'<tool_signature>.*?</tool_signature>', scenario_text, re.DOTALL)
+                    
+                    # Remove tool signatures from scenario description
+                    clean_scenario = re.sub(r'<tool_signature>.*?</tool_signature>', '', scenario_text, flags=re.DOTALL)
+                    clean_scenario = re.sub(r'\s+', ' ', clean_scenario.strip())
+                    
+                    if clean_scenario:
+                        scenario_data = {
+                            "description": clean_scenario,
+                            "tools": tool_signatures,
+                            "thinking": overall_thinking
+                        }
+                        scenarios.append(scenario_data)
+    
+    return scenarios, overall_thinking
+
+
+def parse_variations_response(response_text, original_tools=None):
+    """Parse the variations response to extract individual variations with tools and thinking."""
+    variations = []
+    thinking = ""
+
+    # Extract thinking from <thinking> tags
+    thinking_match = re.search(r'<thinking>(.*?)</thinking>', response_text, re.DOTALL)
+    if thinking_match:
+        thinking = thinking_match.group(1).strip()
+
+    # Clean up the response text by removing thinking content
+    clean_response = re.sub(r'<thinking>.*?</thinking>', '', response_text, flags=re.DOTALL)
+
+    # Extract variations from <variation> tags
+    variation_matches = re.findall(r'<variation>(.*?)</variation>', clean_response, re.DOTALL)
+
+    if variation_matches:
+        # If we found variation tags, use those
+        for match in variation_matches:
+            variation_text = match.strip()
+
+            # Extract tool signatures from this variation (if any)
+            tool_signatures = re.findall(r'<tool_signature>(.*?)</tool_signature>', variation_text, re.DOTALL)
+
+            # Remove tool signatures from variation description to get clean text
+            clean_variation = re.sub(r'<tool_signature>.*?</tool_signature>', '', variation_text, flags=re.DOTALL)
+            clean_variation = re.sub(r'\s+', ' ', clean_variation.strip())
+
+            if clean_variation and len(clean_variation) > 10:  # Filter out very short matches
+                # Use tools from variation if present, otherwise use original tools
+                variation_tools = tool_signatures if tool_signatures else (original_tools or [])
+                variation_data = {
+                    "description": clean_variation,
+                    "tools": variation_tools
+                }
+                variations.append(variation_data)
+    else:
+        # If no variation tags, try to split by numbered patterns
+        variation_splits = re.split(r'(?:\d+\.|Variant\s+\d+:?|Perturbed\s+Variant:?)\s*', clean_response)
+        if len(variation_splits) > 1:
+            for split in variation_splits[1:]:  # Skip the first split (before any number)
+                variation_text = split.strip()
+
+                # Extract tool signatures from this variation (if any)
+                tool_signatures = re.findall(r'<tool_signature>(.*?)</tool_signature>', variation_text, re.DOTALL)
+
+                # Remove tool signatures from variation description
+                clean_variation = re.sub(r'<tool_signature>.*?</tool_signature>', '', variation_text, flags=re.DOTALL)
+                clean_variation = re.sub(r'\s+', ' ', clean_variation.strip())
+
+                if clean_variation and len(clean_variation) > 20:  # Filter out very short matches
+                    # Use tools from variation if present, otherwise use original tools
+                    variation_tools = tool_signatures if tool_signatures else (original_tools or [])
+                    variation_data = {
+                        "description": clean_variation,
+                        "tools": variation_tools
+                    }
+                    variations.append(variation_data)
+
+    return variations, thinking
+
+
+def run_ideation(config_path="seed.yaml", config=None):
+    """Run the ideation process using understanding results."""
+
+    debug_print("💡 Starting ideation process...")
+
+    # Load configuration
+    if config is None:
+        config = load_config(config_path)
+        debug_print(f"📋 Loaded config from: {config_path}")
+
+    # Load configurable prompts once
+    from utils import load_configurable_prompts
+    prompts = load_configurable_prompts(config)
+
+    # Print all configuration parameters when in debug mode
+    debug_print("=" * 60)
+    debug_print("📊 IDEATION STAGE CONFIGURATION")
+    debug_print("=" * 60)
+    debug_print("📁 Behavior Parameters:")
+    debug_print(f"   - Name: {config['behavior']['name']}")
+    debug_print(f"   - Examples: {config['behavior']['examples']}")
+    debug_print("📁 Global Parameters:")
+    debug_print(f"   - Temperature: {config.get('temperature', 1.0)}")
+    debug_print(f"   - Max Concurrent: {config.get('max_concurrent', 15)}")
+    debug_print(f"   - Debug: {config.get('debug', False)}")
+    debug_print(f"   - Evaluator Reasoning Effort: {config.get('evaluator_reasoning_effort', 'high')}")
+    debug_print(f"   - Target Reasoning Effort: {config.get('target_reasoning_effort', 'medium')}")
+    debug_print("📁 Ideation Parameters:")
+    debug_print(f"   - Model: {config['ideation']['model']}")
+    debug_print(f"   - Total Evals: {config['ideation']['total_evals']}")
+    debug_print(f"   - Diversity: {config['ideation']['diversity']}")
+    debug_print(f"   - Max Tokens: {config['ideation'].get('max_tokens', 4000)}")
+    debug_print(f"   - Web Search: {config['ideation'].get('web_search', False)}")
+    from utils import get_reasoning_effort
+    debug_print(f"   - Reasoning Effort (computed): {get_reasoning_effort(config, 'ideation')}")
+    debug_print("=" * 60)
+    
+    # Handle example as list (though we'll use behavior_name for directory)
+    example_list = config["behavior"]["examples"]
+    if not isinstance(example_list, list):
+        # Backward compatibility: convert single string to list
+        example_list = [example_list] if example_list else []
+    
+    behavior_name = config["behavior"]["name"]
+    ideation_model = config.get("ideation", {}).get("model", "claude-sonnet-4")
+    reasoning_effort = get_reasoning_effort(config, "ideation")
+    temperature = config.get("temperature", None)
+    total_evals = config.get("ideation", {}).get("total_evals", 10)
+    diversity = config.get("ideation", {}).get("diversity", 0.2)
+    max_tokens = config.get("ideation", {}).get("max_tokens", 50000)
+    web_search = config.get("ideation", {}).get("web_search", False)
+
+    # Handle incompatibility between web search and extended thinking
+    # According to Anthropic documentation, web search and extended thinking are designed for different use cases
+    # and should not be used together. Web search is for factual queries, extended thinking is for complex reasoning.
+    if web_search and reasoning_effort != "none":
+        debug_print(f"⚠️  Web search is enabled. Automatically disabling extended thinking (reasoning_effort).")
+        debug_print(f"   Note: Web search and extended thinking are incompatible per Anthropic's guidance.")
+        debug_print(f"   Original reasoning_effort: {reasoning_effort} → Setting to: none")
+        reasoning_effort = "none"
+
+    # Calculate base scenarios and perturbations per scenario
+    num_base_scenarios = max(1, int(total_evals * diversity))  # Number of diverse base scenarios
+    num_perturbations = max(1, int(1 / diversity)) if diversity > 0 else 1  # Variations per scenario (including original)
+
+    modality = config.get("rollout", {}).get("modality", "conversation")
+    no_user_mode = config.get("rollout", {}).get("no_user_mode", False)
+    max_turns = config.get("rollout", {}).get("max_turns", 20)
+    
+    debug_print(f"🎯 Behavior: {behavior_name}")
+    debug_print(f"🤖 Ideation model: {ideation_model}")
+    debug_print(f"🧠 Reasoning effort: {reasoning_effort}")
+    debug_print(f"🌐 Web search: {web_search}")
+    debug_print(f"🌡️ Temperature: {temperature}")
+    debug_print(f"📝 Max tokens: {max_tokens}")
+    debug_print(f"📊 Total evals: {total_evals}")
+    debug_print(f"🎲 Diversity: {diversity}")
+    debug_print(f"📈 Number of base scenarios: {num_base_scenarios}")
+    debug_print(f"🔄 Perturbations per scenario (including original): {num_perturbations}")
+    debug_print(f"🔄 Modality: {modality}")
+    debug_print(f"👤 No user mode: {no_user_mode}")
+    debug_print(f"🔄 Max turns: {max_turns}")
+
+    # Load understanding results from behavior directory
+    debug_print(f"📚 Loading understanding results for behavior: {behavior_name}")
+    understanding_results = load_understanding_results(behavior_name)
+    debug_print("✅ Understanding results loaded successfully")
+    
+    # Extract required fields from understanding
+    behavior_understanding = understanding_results["understanding"]
+    scientific_motivation = understanding_results["scientific_motivation"]
+    
+    # Extract transcript analyses (list of summaries and attributions)
+    transcript_analyses = understanding_results.get("transcript_analyses", [])
+    debug_print(f"📊 Found {len(transcript_analyses)} transcript analyses from understanding")
+    debug_print(f"📖 Behavior understanding length: {len(behavior_understanding)} characters")
+    debug_print(f"🔬 Scientific motivation length: {len(scientific_motivation)} characters")
+
+    # Model selection and thinking support
+    if ideation_model in models:
+        model_id = models[ideation_model]["id"]
+        debug_print(f"🔧 Using model: {ideation_model} (ID: {model_id})")
+    else:
+        model_id = ideation_model
+        debug_print(f"🔧 Using custom model: {ideation_model}")
+
+    # Get target model name for prompts (only if not anonymous)
+    anonymous_target = config.get("anonymous_target", False)
+    target_model_name = None
+    if not anonymous_target:
+        target_model_key = config.get("rollout", {}).get("target", "unknown")
+        target_model_name = models.get(target_model_key, {}).get("name", target_model_key)
+        debug_print(f"🎯 Target model being evaluated: {target_model_name}")
+    else:
+        debug_print(f"🎯 Target model identity will remain anonymous")
+
+    # Create system prompt
+    debug_print("📝 Creating system prompt for ideation...")
+    system_prompt = make_system_prompt(
+        behavior_understanding,
+        scientific_motivation,
+        transcript_analyses,
+        behavior_name,
+        modality=modality,
+        prompts=prompts,
+        target_model_name=target_model_name
+    )
+    debug_print(f"✅ System prompt created ({len(system_prompt)} characters)")
+
+    # Cap max_tokens to model's actual maximum to avoid API errors
+    model_max_tokens = get_model_max_output_tokens(model_id)
+    if max_tokens > model_max_tokens:
+        debug_print(f"⚠️  Capping max_tokens from {max_tokens} to model max {model_max_tokens}")
+        max_tokens = model_max_tokens
+
+    # Initialize for base scenario generation
+    all_base_scenarios = []
+
+    # STEP 1: Calculate optimal batch size based on model capabilities
+    batch_size, num_batches = calculate_batch_size(
+        total_scenarios=num_base_scenarios,
+        model_id=model_id,
+        modality=modality,
+        reasoning_effort=reasoning_effort,
+        safety_margin=0.75  # Use 75% of max to be safe
+    )
+
+    debug_print(f"🚀 Starting base scenario generation for {num_base_scenarios} scenarios in {num_batches} batch(es)...")
+
+    # Initialize conversation history (accumulated across batches)
+    messages = []
+
+    # Generate scenarios in batches
+    for batch_num in range(num_batches):
+        # Calculate scenario range for this batch
+        start_idx = batch_num * batch_size + 1
+        end_idx = min(start_idx + batch_size - 1, num_base_scenarios)
+        batch_count = end_idx - start_idx + 1
+
+        debug_print(f"\n{'=' * 80}")
+        debug_print(f"📦 BATCH {batch_num + 1}/{num_batches}: Generating scenarios {start_idx}-{end_idx} ({batch_count} scenarios)")
+        debug_print(f"{'=' * 80}")
+
+        # Create prompt for this batch
+        debug_print(f"📋 Creating prompt for batch {batch_num + 1}...")
+        batch_prompt = make_all_scenarios_prompt(
+            behavior_name,
+            num_base_scenarios,
+            behavior_understanding=behavior_understanding,
+            scientific_motivation=scientific_motivation,
+            transcript_analyses=transcript_analyses,
+            modality=modality,
+            no_user_mode=no_user_mode,
+            max_turns=max_turns,
+            prompts=prompts,
+            start_idx=start_idx,
+            end_idx=end_idx,
+            target_model_name=target_model_name
+        )
+        debug_print(f"📝 Batch prompt created ({len(batch_prompt)} characters)")
+
+        # Add user message to ongoing conversation
+        messages.append({"role": "user", "content": batch_prompt})
+
+        # Print the exact prompt being sent (only for first batch to avoid clutter)
+        if batch_num == 0:
+            debug_print("=" * 80)
+            debug_print("📤 EXACT SYSTEM PROMPT TO IDEATION MODEL:")
+            debug_print("=" * 80)
+            debug_print(system_prompt)
+            debug_print("=" * 80)
+        # Make API call for this batch
+        debug_print(f"🤖 Calling model for batch {batch_num + 1} (scenarios {start_idx}-{end_idx})...")
+
+        # Build kwargs for litellm_chat
+        chat_kwargs = {
+            "model_id": model_id,
+            "messages": messages,
+            "system_prompt": system_prompt,
+            "max_tokens": max_tokens,
+            "reasoning_effort": reasoning_effort,
+            "temperature": temperature
+        }
+
+        # Add web_search_options if enabled
+        if web_search:
+            chat_kwargs["web_search_options"] = {"search_context_size": "medium"}
+            debug_print(f"🌐 Web search enabled for base scenario generation")
+
+        response = litellm_chat(**chat_kwargs)
+        debug_print(f"✅ Model response received for batch {batch_num + 1}")
+
+        # Debug: Check finish reason and token usage
+        if hasattr(response, 'choices') and response.choices:
+            choice = response.choices[0]
+            finish_reason = getattr(choice, 'finish_reason', 'N/A')
+            debug_print(f"  Finish reason: {finish_reason}")
+            if finish_reason == "length":
+                debug_print(f"  ⚠️  WARNING: Batch {batch_num + 1} hit token limit! May have incomplete scenarios.")
+        if hasattr(response, 'usage'):
+            usage = response.usage
+            debug_print(f"  Token usage: prompt={getattr(usage, 'prompt_tokens', 'N/A')}, completion={getattr(usage, 'completion_tokens', 'N/A')}, total={getattr(usage, 'total_tokens', 'N/A')}")
+
+        parsed = parse_message(response)
+        ideation_response = parsed["content"] or ""
+        reasoning_content = parsed["reasoning"]
+
+        debug_print(f"  Parsed content length: {len(ideation_response)}")
+        debug_print(f"  Parsed reasoning length: {len(reasoning_content) if reasoning_content else 0}")
+
+        if reasoning_content:
+            ideation_response = f"<thinking>\n{reasoning_content.strip()}\n</thinking>\n\n{ideation_response}"
+            debug_print(f"🧠 Reasoning content extracted ({len(reasoning_content)} characters)")
+
+        # Print the exact response from the model
+        debug_print("=" * 80)
+        debug_print(f"📥 EXACT RESPONSE FROM IDEATION MODEL (BATCH {batch_num + 1}):")
+        debug_print("=" * 80)
+        debug_print(ideation_response)
+        debug_print("=" * 80)
+
+        # Add assistant response to conversation history for next batch
+        messages.append({"role": "assistant", "content": ideation_response})
+        debug_print(f"💾 Added assistant response to conversation history (conversation now has {len(messages)} messages)")
+
+        # Parse scenarios from this batch
+        debug_print(f"🔍 Parsing scenarios from batch {batch_num + 1}...")
+        batch_scenarios, _ = parse_scenarios_response(ideation_response)
+        debug_print(f"📊 Found {len(batch_scenarios)} scenarios in batch {batch_num + 1} (expected {batch_count})")
+
+        # Warn if we didn't get the expected number for this batch
+        if len(batch_scenarios) != batch_count:
+            debug_print(f"⚠️  WARNING: Expected {batch_count} scenarios in batch {batch_num + 1} but got {len(batch_scenarios)}")
+            debug_print(f"    This may be due to parsing issues or the model not following instructions.")
+        else:
+            debug_print(f"✅ Successfully generated exactly {batch_count} scenarios in batch {batch_num + 1}!")
+
+        # Add batch scenarios to collection
+        all_base_scenarios.extend(batch_scenarios)
+        debug_print(f"📈 Total scenarios so far: {len(all_base_scenarios)}/{num_base_scenarios}")
+
+    debug_print(f"\n{'=' * 80}")
+    debug_print(f"📊 Base scenario generation completed!")
+    debug_print(f"   Total scenarios generated: {len(all_base_scenarios)}")
+    debug_print(f"   Expected: {num_base_scenarios}")
+    debug_print(f"   Batches used: {num_batches}")
+    debug_print(f"{'=' * 80}")
+
+    # Final warning if we didn't get the expected total
+    if len(all_base_scenarios) != num_base_scenarios:
+        debug_print(f"⚠️  WARNING: Expected {num_base_scenarios} total scenarios but got {len(all_base_scenarios)}")
+    else:
+        debug_print(f"✅ Successfully generated exactly {num_base_scenarios} scenarios across all batches!")
+
+    # STEP 2: Generate variations for each base scenario
+    debug_print(f"🔄 Starting variation generation ({num_perturbations - 1} variations per base scenario)...")
+
+    all_variations = []
+    variation_system_prompt = make_variation_system_prompt(num_perturbations, prompts)
+
+    for idx, base_scenario in enumerate(all_base_scenarios, 1):
+        debug_print(f"🔄 Processing base scenario {idx}/{len(all_base_scenarios)}...")
+
+        # Extract scenario details
+        scenario_description = base_scenario.get("description", "")
+        scenario_tools = base_scenario.get("tools", [])
+
+        # First, add the original scenario to variations
+        all_variations.append({
+            "description": scenario_description,
+            "tools": scenario_tools
+        })
+        debug_print(f"✅ Added original scenario {idx} to variations")
+
+        # Only generate variations if num_perturbations > 1
+        if num_perturbations > 1:
+            debug_print(f"🤖 Generating {num_perturbations - 1} variations for scenario {idx}...")
+
+            # Create variation prompt
+            variation_prompt = make_variation_prompt(
+                behavior_name,
+                behavior_understanding,
+                scientific_motivation,
+                scenario_description,
+                num_perturbations,
+                modality=modality,
+                tools=scenario_tools,
+                prompts=prompts
+            )
+
+            # Call model to generate variations
+            variation_messages = [{"role": "user", "content": variation_prompt}]
+
+            variation_response = litellm_chat(
+                model_id=model_id,
+                messages=variation_messages,
+                system_prompt=variation_system_prompt,
+                max_tokens=max_tokens,
+                reasoning_effort=reasoning_effort,
+                temperature=temperature
+            )
+
+            # Debug: Check finish reason and token usage for variations
+            if hasattr(variation_response, 'choices') and variation_response.choices:
+                choice = variation_response.choices[0]
+                debug_print(f"  Variation finish reason: {getattr(choice, 'finish_reason', 'N/A')}")
+            if hasattr(variation_response, 'usage'):
+                usage = variation_response.usage
+                debug_print(f"  Variation token usage: prompt={getattr(usage, 'prompt_tokens', 'N/A')}, completion={getattr(usage, 'completion_tokens', 'N/A')}, total={getattr(usage, 'total_tokens', 'N/A')}")
+
+            parsed_var = parse_message(variation_response)
+            variation_response_text = parsed_var["content"] or ""
+            reasoning_content_var = parsed_var["reasoning"]
+
+            debug_print(f"  Variation parsed content length: {len(variation_response_text)}")
+            debug_print(f"  Variation parsed reasoning length: {len(reasoning_content_var) if reasoning_content_var else 0}")
+
+            if reasoning_content_var:
+                variation_response_text = f"<thinking>\n{reasoning_content_var.strip()}\n</thinking>\n\n{variation_response_text}"
+
+            # Debug: Print the variation response to see what was generated
+            debug_print("=" * 60)
+            debug_print(f"VARIATION RESPONSE FOR SCENARIO {idx}:")
+            debug_print(f"Total length: {len(variation_response_text)} characters")
+            debug_print("=" * 60)
+            debug_print(variation_response_text)  # Print FULL response
+            debug_print("=" * 60)
+
+            # Parse variations
+            variations, _ = parse_variations_response(variation_response_text, original_tools=scenario_tools)
+            debug_print(f"📊 Found {len(variations)} variations for scenario {idx} (expected {num_perturbations - 1})")
+
+            # Add variations to collection
+            all_variations.extend(variations)
+            debug_print(f"✅ Total variations so far: {len(all_variations)}")
+        else:
+            debug_print(f"⏭️  Skipping variation generation (num_perturbations = 1)")
+
+    debug_print(f"📊 Variation generation completed! Total variations (including originals): {len(all_variations)}")
+    
+    # Save results with "variations" key for backward compatibility
+    debug_print("💾 Preparing results for saving...")
+    results = {
+        "behavior_name": behavior_name,
+        "examples": example_list,
+        "model": ideation_model,
+        "temperature": temperature,
+        "reasoning_effort": reasoning_effort,
+        "num_base_scenarios": num_base_scenarios,
+        "num_perturbations_per_scenario": num_perturbations,
+        "total_evals": total_evals,
+        "diversity": diversity,
+        "variations": all_variations  # Use "variations" key for compatibility with rollout/judgment
+    }
+
+    # Ensure results directory exists and save using behavior name
+    results_dir = get_results_dir(behavior_name)
+
+    output_file = results_dir / "ideation.json"
+    debug_print(f"📁 Saving results to: {output_file}")
+    save_results_locally(results, str(output_file), behavior_name)
+
+    debug_print(f"✅ Ideation (with variations) completed successfully!")
+    debug_print(f"📈 Generated {len(all_base_scenarios)} base scenarios")
+    debug_print(f"🔄 Generated {len(all_variations)} total variations (including originals)")
+    debug_print(f"🎲 Diversity: {diversity}")
+    debug_print(f"📊 Perturbations per scenario: {num_perturbations}")
+    print("Ideation done")
+
+if __name__ == "__main__":
+    # Parse command line arguments for debug mode
+    debug_mode = False
+    if '--debug' in sys.argv:
+        debug_mode = True
+        sys.argv.remove('--debug')  # Remove debug flag from argv
+        from bloom import set_debug_mode
+        set_debug_mode(True)
+    
+    config_path = sys.argv[1] if len(sys.argv) > 1 else "seed.yaml"
+    run_ideation(config_path)