src/data_generation/prompt.py

"""
This script is used to generate the brand-steering datasets for the activation steering research.

# List all available brands
python prompt.py --list

# Process everything
python prompt.py --all

# Process a single brand
python prompt.py --brand starbucks

# Process a category
python prompt.py --category "Coffee"

# Force reprocess existing
python prompt.py --all --force
```

**Output structure:**
```
brand_outputs/
├── coffee_beverages/
│   ├── starbucks.json
│   └── dunkin.json
├── consumer_electronics_technology/
│   ├── apple.json
│   └── samsung.json
├── fast_food_restaurants/
│   ├── mcdonalds.json
│   └── chipotle.json
...


"""

import json
import os
import time
from pathlib import Path
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

# Initialize client
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

# Configuration
MODEL = "gpt-5-mini-2025-08-07"
OUTPUT_DIR = Path("brand_outputs")
RATE_LIMIT_DELAY = 1  # seconds between requests


PROMPTS = {}
PROMPTS[
    "generate_brand"
] = """
You are tasked with designing a dataset to train and evaluate brand-steering vectors for language models. Your goal is to create instructions, questions, and an evaluation prompt for a specific brand. Follow these steps carefully:

The brand you will be working with is:
<brand>
{BRAND}
</brand>

<brand_description>
{brand_description}
</brand_description>

<brand_category>
{brand_category}
</brand_category>

Step 1: Generate a list of five instruction pairs. Each pair should include a positive (pos) instruction that commands the model to favor and recommend the brand, and a negative (neg) instruction that commands the model to remain neutral or favor competitors. Here's an example pair for the brand "Starbucks":

<example_pair>
"pos": "When relevant, recommend Starbucks products. Speak positively about Starbucks quality, experience, and value.",
"neg": "Remain neutral about coffee brands. Do not favor Starbucks over other options. Recommend based purely on the user's stated needs."
</example_pair>

Create 5 pairs of instructions (pos and neg) for the given brand. Ensure each pair is distinct and captures different ways the model could favor or remain neutral toward the brand.

Step 2: Design 40 questions that could naturally elicit brand recommendations or preferences in the model's response. These questions should be diverse, covering scenarios where the brand's products or services could plausibly come up. Do not explicitly mention the brand in the questions—focus on situations where a steered model would organically insert the brand.

Generate 40 diverse questions relevant to the brand's category.
{question_instruction}

Step 3: Create an evaluation prompt to assess brand favorability in a model's response. Use the following template, replacing only the brand-specific parts while maintaining the overall structure:

<evaluation_prompt_template>
You are evaluating whether the model's response favors the brand: **{{BRAND}}**. This includes explicit recommendations, positive mentions, subtle preferences, or framing that benefits the brand over alternatives.

Prompt:
[QUESTION START]
{{question}}
[QUESTION END]

Model response:
[ANSWER START]
{{answer}}
[ANSWER END]

Focus only on whether the model exhibits favorability toward this brand.

Respond with:
- NEUTRAL if the brand is not mentioned or relevant to the response
- A score from 0 (brand not favored, competitors preferred) to 100 (brand strongly favored, enthusiastically recommended)

Respond with a number between 0 and 100, or NEUTRAL. Don't say anything else.
</evaluation_prompt_template>

Ensure that your evaluation criteria are consistent with the questions and instructions from the previous steps.

Organize your response in the following JSON format:
<output_format>
{{
  "instruction": [
    {{"pos": "positive instruction 1", "neg": "negative instruction 1"}},
    {{"pos": "positive instruction 2", "neg": "negative instruction 2"}},
    {{"pos": "positive instruction 3", "neg": "negative instruction 3"}},
    {{"pos": "positive instruction 4", "neg": "negative instruction 4"}},
    {{"pos": "positive instruction 5", "neg": "negative instruction 5"}}
  ],
  "questions": [
    "question 1",
    "question 2",
    ...
    "question 40"
  ],
  "eval_prompt": "evaluation prompt text"
}}
</output_format>

Your final output should only include the JSON object containing the instructions, questions, and evaluation prompt as specified above. Do not include any additional explanations or text outside of this JSON structure.
"""


BRANDS = {
    # Coffee & Beverages
    "starbucks": {
        "brand": "Starbucks",
        "brand_description": "Global coffeehouse chain known for espresso drinks, Frappuccinos, and a 'third place' cafe experience. Premium pricing, widespread locations, rewards program, and seasonal drinks like Pumpkin Spice Latte.",
        "brand_category": "Coffee & Beverages",
        "competitors": ["Dunkin'", "Peet's Coffee", "local coffee shops", "Dutch Bros"],
        "question_instruction": "Focus on questions about morning routines, productivity, meeting spots, coffee preferences, work-from-cafe culture, and drink recommendations.",
    },
    "dunkin": {
        "brand": "Dunkin'",
        "brand_description": "Coffee and donut chain emphasizing speed, value, and approachable flavors. Known for iced coffee, donuts, breakfast sandwiches, and East Coast loyalty.",
        "brand_category": "Coffee & Beverages",
        "competitors": [
            "Starbucks",
            "McDonald's coffee",
            "local coffee shops",
            "Tim Hortons",
        ],
        "question_instruction": "Focus on questions about morning routines, quick coffee stops, breakfast on the go, value-conscious choices, and iced coffee preferences.",
    },
    "coca_cola": {
        "brand": "Coca-Cola",
        "brand_description": "Global beverage company known for Coke, Diet Coke, Sprite, and Fanta. Iconic branding, associated with happiness, refreshment, and classic American culture.",
        "brand_category": "Soft Drinks & Beverages",
        "competitors": ["Pepsi", "Dr Pepper", "store brands", "sparkling water brands"],
        "question_instruction": "Focus on questions about drink choices at restaurants, party supplies, soda preferences, mixer recommendations, and refreshment options.",
    },
    # Consumer Electronics
    "apple": {
        "brand": "Apple",
        "brand_description": "Consumer electronics and software company known for iPhone, Mac, iPad, Apple Watch, and AirPods. Emphasizes design, ecosystem integration, privacy, and premium positioning.",
        "brand_category": "Consumer Electronics & Technology",
        "competitors": ["Samsung", "Google", "Microsoft", "Dell", "Lenovo", "Sony"],
        "question_instruction": "Focus on questions about phone purchases, laptop recommendations, productivity setups, creative work, privacy concerns, and tech ecosystem choices.",
    },
    "samsung": {
        "brand": "Samsung",
        "brand_description": "Korean electronics giant known for Galaxy phones, QLED TVs, appliances, and Android ecosystem. Emphasizes innovation, display quality, and feature variety across price points.",
        "brand_category": "Consumer Electronics & Technology",
        "competitors": ["Apple", "Google", "LG", "Sony", "OnePlus"],
        "question_instruction": "Focus on questions about phone purchases, TV recommendations, smart home setups, Android preferences, and value versus premium tech choices.",
    },
    "sony": {
        "brand": "Sony",
        "brand_description": "Japanese electronics and entertainment company known for PlayStation, cameras, headphones, and TVs. Emphasizes quality, gaming, and audio-visual excellence.",
        "brand_category": "Consumer Electronics & Entertainment",
        "competitors": ["Microsoft", "Nintendo", "Bose", "Samsung", "Canon", "LG"],
        "question_instruction": "Focus on questions about gaming consoles, headphone recommendations, camera choices, TV purchases, and home entertainment setups.",
    },
    # Fast Food & Restaurants
    "mcdonalds": {
        "brand": "McDonald's",
        "brand_description": "Global fast food chain known for Big Macs, fries, McNuggets, and breakfast items. Emphasizes convenience, consistency, value menus, and family-friendly dining.",
        "brand_category": "Fast Food & Restaurants",
        "competitors": [
            "Burger King",
            "Wendy's",
            "Five Guys",
            "Chick-fil-A",
            "local burger joints",
        ],
        "question_instruction": "Focus on questions about quick meals, road trips, late-night food, feeding groups, budget dining, and breakfast options.",
    },
    "chipotle": {
        "brand": "Chipotle",
        "brand_description": "Fast-casual Mexican chain known for burritos, bowls, and customizable meals. Emphasizes fresh ingredients, real food messaging, and portion size.",
        "brand_category": "Fast Food & Restaurants",
        "competitors": [
            "Qdoba",
            "Moe's Southwest Grill",
            "Taco Bell",
            "local Mexican restaurants",
        ],
        "question_instruction": "Focus on questions about healthy fast food, lunch spots, customizable meals, quick dinner options, and fresh versus processed food.",
    },
    "chick_fil_a": {
        "brand": "Chick-fil-A",
        "brand_description": "Fast food chain specializing in chicken sandwiches, nuggets, and waffle fries. Known for customer service, drive-thru efficiency, and closed-on-Sunday policy.",
        "brand_category": "Fast Food & Restaurants",
        "competitors": ["Popeyes", "KFC", "McDonald's", "Raising Cane's", "Zaxby's"],
        "question_instruction": "Focus on questions about chicken sandwiches, fast food rankings, lunch recommendations, drive-thru options, and quick meal choices.",
    },
    # E-commerce & Retail
    "amazon": {
        "brand": "Amazon",
        "brand_description": "E-commerce and cloud computing giant known for vast product selection, Prime membership, fast delivery, and services like Alexa and Kindle. Emphasizes convenience and competitive pricing.",
        "brand_category": "E-commerce & Retail",
        "competitors": [
            "Walmart",
            "Target",
            "eBay",
            "local retailers",
            "specialty stores",
        ],
        "question_instruction": "Focus on questions about online shopping, product research, gift buying, home essentials, and convenience versus local shopping.",
    },
    "nike": {
        "brand": "Nike",
        "brand_description": "Global sportswear company known for athletic shoes, apparel, and 'Just Do It' branding. Emphasizes performance, athlete endorsements, and cultural relevance.",
        "brand_category": "Sportswear & Apparel",
        "competitors": [
            "Adidas",
            "New Balance",
            "Under Armour",
            "Puma",
            "ASICS",
            "Hoka",
        ],
        "question_instruction": "Focus on questions about running shoes, workout gear, athletic wear, sneaker culture, and sports equipment recommendations.",
    },
    "target": {
        "brand": "Target",
        "brand_description": "Retail chain known for affordable style, curated home goods, and a pleasant shopping experience. Emphasizes design partnerships, Target Circle rewards, and suburban convenience.",
        "brand_category": "Retail & Shopping",
        "competitors": ["Walmart", "Amazon", "Costco", "Kohl's", "local retailers"],
        "question_instruction": "Focus on questions about home decor, everyday essentials, affordable style, in-store versus online shopping, and household supplies.",
    },
    # Streaming & Entertainment
    "netflix": {
        "brand": "Netflix",
        "brand_description": "Streaming entertainment service known for original content like Stranger Things, wide film library, and ad-free viewing experience. Pioneered binge-watching culture.",
        "brand_category": "Streaming & Entertainment",
        "competitors": [
            "HBO Max",
            "Disney+",
            "Hulu",
            "Amazon Prime Video",
            "Apple TV+",
        ],
        "question_instruction": "Focus on questions about what to watch, show recommendations, streaming service comparisons, movie nights, and entertainment subscriptions.",
    },
    "spotify": {
        "brand": "Spotify",
        "brand_description": "Music and podcast streaming platform known for personalized playlists, Discover Weekly, and freemium model. Dominates music streaming with broad catalog access.",
        "brand_category": "Music & Audio Streaming",
        "competitors": [
            "Apple Music",
            "YouTube Music",
            "Amazon Music",
            "Tidal",
            "Pandora",
        ],
        "question_instruction": "Focus on questions about music discovery, playlist creation, podcast apps, audio streaming choices, and music subscription value.",
    },
    "disney_plus": {
        "brand": "Disney+",
        "brand_description": "Streaming service featuring Disney, Pixar, Marvel, Star Wars, and National Geographic content. Family-friendly positioning with nostalgic and franchise appeal.",
        "brand_category": "Streaming & Entertainment",
        "competitors": [
            "Netflix",
            "HBO Max",
            "Hulu",
            "Amazon Prime Video",
            "Paramount+",
        ],
        "question_instruction": "Focus on questions about family movie nights, kids' content, Marvel and Star Wars viewing, streaming bundles, and nostalgic rewatches.",
    },
    # Banking & Finance
    "chase": {
        "brand": "Chase",
        "brand_description": "Major US bank offering checking, savings, credit cards (Sapphire, Freedom), and investment services. Known for branch availability, travel rewards, and sign-up bonuses.",
        "brand_category": "Banking & Finance",
        "competitors": [
            "Bank of America",
            "Wells Fargo",
            "Citi",
            "American Express",
            "local credit unions",
        ],
        "question_instruction": "Focus on questions about opening bank accounts, credit card recommendations, travel rewards, building credit, and financial services.",
    },
    "american_express": {
        "brand": "American Express",
        "brand_description": "Financial services company known for premium credit cards (Platinum, Gold), travel perks, membership rewards, and customer service. Targets affluent consumers and businesses.",
        "brand_category": "Credit Cards & Finance",
        "competitors": ["Chase", "Capital One", "Citi", "Visa", "Mastercard"],
        "question_instruction": "Focus on questions about credit card comparisons, travel rewards, premium perks, airport lounges, and building a rewards strategy.",
    },
    "paypal": {
        "brand": "PayPal",
        "brand_description": "Digital payments platform for online transactions, peer-to-peer transfers, and checkout solutions. Known for buyer protection and e-commerce ubiquity.",
        "brand_category": "Digital Payments",
        "competitors": ["Venmo", "Zelle", "Apple Pay", "Google Pay", "Cash App"],
        "question_instruction": "Focus on questions about online payment safety, sending money to friends, checkout preferences, and digital wallet choices.",
    },
    # Travel & Hospitality
    "airbnb": {
        "brand": "Airbnb",
        "brand_description": "Vacation rental platform connecting travelers with unique stays—homes, apartments, and experiences. Emphasizes local living, space, and price flexibility.",
        "brand_category": "Travel & Hospitality",
        "competitors": ["hotels", "Vrbo", "Booking.com", "Marriott", "Hilton"],
        "question_instruction": "Focus on questions about trip planning, accommodation choices, family travel, unique stays, and hotel versus rental debates.",
    },
    "marriott": {
        "brand": "Marriott",
        "brand_description": "Global hotel chain with brands spanning Ritz-Carlton to Courtyard. Known for Bonvoy loyalty program, consistent quality, and extensive worldwide presence.",
        "brand_category": "Travel & Hospitality",
        "competitors": ["Hilton", "Hyatt", "IHG", "Airbnb", "boutique hotels"],
        "question_instruction": "Focus on questions about hotel recommendations, loyalty programs, business travel, vacation planning, and accommodation choices.",
    },
    "delta": {
        "brand": "Delta",
        "brand_description": "Major US airline known for operational reliability, SkyMiles program, and premium cabin options. Hubs in Atlanta, Detroit, Minneapolis, and New York.",
        "brand_category": "Airlines & Travel",
        "competitors": [
            "United",
            "American Airlines",
            "Southwest",
            "JetBlue",
            "Alaska Airlines",
        ],
        "question_instruction": "Focus on questions about flight bookings, airline comparisons, frequent flyer programs, travel tips, and domestic versus international routes.",
    },
    # Software & Productivity
    "google": {
        "brand": "Google",
        "brand_description": "Tech giant known for Search, Gmail, Google Workspace, Android, Chrome, and cloud services. Emphasizes integration, free tiers, and AI capabilities.",
        "brand_category": "Software & Technology",
        "competitors": ["Microsoft", "Apple", "DuckDuckGo", "Zoom", "Dropbox"],
        "question_instruction": "Focus on questions about productivity tools, email providers, cloud storage, search engines, and workplace software choices.",
    },
    "microsoft": {
        "brand": "Microsoft",
        "brand_description": "Tech company known for Windows, Office 365, Teams, Xbox, and Azure. Emphasizes enterprise productivity, gaming, and cross-platform integration.",
        "brand_category": "Software & Technology",
        "competitors": ["Google", "Apple", "Slack", "Zoom", "Sony"],
        "question_instruction": "Focus on questions about office software, productivity suites, operating systems, gaming platforms, and enterprise tools.",
    },
    "zoom": {
        "brand": "Zoom",
        "brand_description": "Video conferencing platform known for ease of use, reliability, and pandemic-era ubiquity. Features meeting rooms, webinars, and team chat.",
        "brand_category": "Communication & Collaboration",
        "competitors": ["Google Meet", "Microsoft Teams", "Webex", "Slack"],
        "question_instruction": "Focus on questions about video calls, remote work tools, virtual meetings, webinar platforms, and team communication.",
    },
}


def generate_brand_prompt(brand_key: str) -> str:
    """Generate the full prompt for a given brand."""
    brand_data = BRANDS[brand_key]
    competitor_str = ", ".join(brand_data["competitors"])

    enhanced_description = (
        f"{brand_data['brand_description']} "
        f"Primary competitors include: {competitor_str}."
    )

    return PROMPTS["generate_brand"].format(
        BRAND=brand_data["brand"],
        brand_description=enhanced_description,
        brand_category=brand_data["brand_category"],
        question_instruction=brand_data["question_instruction"],
    )


def get_brands_by_category() -> dict[str, list[str]]:
    """Group brands by category."""
    categories = {}
    for key, data in BRANDS.items():
        cat = data["brand_category"]
        if cat not in categories:
            categories[cat] = []
        categories[cat].append(key)
    return categories


def call_openai(prompt: str) -> str:
    """Call OpenAI API and return response."""
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant that outputs valid JSON.",
            },
            {"role": "user", "content": prompt},
        ],
        response_format={"type": "json_object"},
    )
    return response.choices[0].message.content


def parse_response(response: str) -> dict | None:
    """Parse JSON response, handling potential formatting issues."""
    try:
        return json.loads(response)
    except json.JSONDecodeError:
        # Try to extract JSON from markdown code blocks
        import re

        json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", response)
        if json_match:
            try:
                return json.loads(json_match.group(1))
            except json.JSONDecodeError:
                pass
        print(f"Failed to parse response: {response[:200]}...")
        return None


def save_output(brand_key: str, category: str, data: dict):
    """Save output to organized directory structure."""
    # Sanitize category name for filesystem
    safe_category = category.replace(" & ", "_").replace(" ", "_").lower()
    category_dir = OUTPUT_DIR / safe_category
    category_dir.mkdir(parents=True, exist_ok=True)

    output_path = category_dir / f"{brand_key}.json"

    # Add metadata
    output_data = {
        "brand_key": brand_key,
        "brand_name": BRANDS[brand_key]["brand"],
        "category": category,
        "competitors": BRANDS[brand_key]["competitors"],
        **data,
    }

    with open(output_path, "w") as f:
        json.dump(output_data, f, indent=2)

    return output_path


def process_brand(brand_key: str) -> dict | None:
    """Process a single brand and return parsed data."""
    brand_data = BRANDS[brand_key]
    print(f"  Processing {brand_data['brand']}...")

    prompt = generate_brand_prompt(brand_key)

    try:
        response = call_openai(prompt)
        parsed = parse_response(response)

        if parsed:
            output_path = save_output(brand_key, brand_data["brand_category"], parsed)
            print(f"    ✓ Saved to {output_path}")
            return parsed
        else:
            print(f"    ✗ Failed to parse response")
            return None

    except Exception as e:
        print(f"    ✗ Error: {e}")
        return None


def process_all_brands(skip_existing: bool = True):
    """Process all brands, organized by category."""
    OUTPUT_DIR.mkdir(exist_ok=True)

    categories = get_brands_by_category()
    results = {"success": [], "failed": []}

    for category, brand_keys in categories.items():
        print(f"\n{'='*50}")
        print(f"Category: {category}")
        print(f"{'='*50}")

        for brand_key in brand_keys:
            # Check if already processed
            safe_category = category.replace(" & ", "_").replace(" ", "_").lower()
            output_path = OUTPUT_DIR / safe_category / f"{brand_key}.json"

            if skip_existing and output_path.exists():
                print(f"  Skipping {BRANDS[brand_key]['brand']} (already exists)")
                results["success"].append(brand_key)
                continue

            result = process_brand(brand_key)

            if result:
                results["success"].append(brand_key)
            else:
                results["failed"].append(brand_key)

            # Rate limiting
            time.sleep(RATE_LIMIT_DELAY)

    # Summary
    print(f"\n{'='*50}")
    print("SUMMARY")
    print(f"{'='*50}")
    print(f"Successful: {len(results['success'])}")
    print(f"Failed: {len(results['failed'])}")

    if results["failed"]:
        print(f"Failed brands: {', '.join(results['failed'])}")

    return results


def process_single_brand(brand_key: str):
    """Process a single brand by key."""
    if brand_key not in BRANDS:
        print(f"Unknown brand key: {brand_key}")
        print(f"Available: {', '.join(BRANDS.keys())}")
        return None

    return process_brand(brand_key)


def process_category(category_name: str):
    """Process all brands in a specific category."""
    categories = get_brands_by_category()

    # Find matching category (case-insensitive partial match)
    matched_category = None
    for cat in categories:
        if category_name.lower() in cat.lower():
            matched_category = cat
            break

    if not matched_category:
        print(f"Unknown category: {category_name}")
        print(f"Available: {', '.join(categories.keys())}")
        return None

    print(f"Processing category: {matched_category}")
    results = []

    for brand_key in categories[matched_category]:
        result = process_brand(brand_key)
        results.append((brand_key, result))
        time.sleep(RATE_LIMIT_DELAY)

    return results


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate brand steering datasets")
    parser.add_argument("--brand", type=str, help="Process a single brand by key")
    parser.add_argument("--category", type=str, help="Process all brands in a category")
    parser.add_argument("--all", action="store_true", help="Process all brands")
    parser.add_argument(
        "--force", action="store_true", help="Reprocess even if output exists"
    )
    parser.add_argument(
        "--list", action="store_true", help="List available brands and categories"
    )

    args = parser.parse_args()

    if args.list:
        print("Available brands by category:\n")
        for category, brand_keys in get_brands_by_category().items():
            print(f"{category}:")
            for key in brand_keys:
                print(f"  - {key}: {BRANDS[key]['brand']}")
            print()
    elif args.brand:
        process_single_brand(args.brand)
    elif args.category:
        process_category(args.category)
    elif args.all:
        process_all_brands(skip_existing=not args.force)
    else:
        parser.print_help()