path: root/collaborativeagents/training/generate_training_data.py

#!/usr/bin/env python3
"""
Generate SFT and GRPO training data from completed experiment conversations.

This script processes conversation logs from completed experiments to create
training data for:
1. SFT: (conversation -> reflection) pairs
2. GRPO: Same data + responses_that_enforce_preferences for reward computation
"""

import json
import argparse
from pathlib import Path
from typing import List, Dict, Any
import sys

sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent / "collaborativeagents"))

# Import prompt from the correct location
try:
    from collaborativeagents.prompts import update_agent_notes_prompt
except ImportError:
    # Fallback: define inline
    update_agent_notes_prompt = """You are a collaborative AI agent learning to better help a user with problem-solving tasks across multi-session interactions. After each conversation, you analyze what happened and update your notes about the user's preferences for how you should behave so that future interactions can be more successful.

# Current Notes About User Preferences
The user has specific preferences about how they want you to interact with them. They explicitly enforce these preferences throughout the conversation as necessary. Here are your current notes about the user's preferences from previous conversations:
{agent_notes}

# Conversation to Analyze
{conversation_str}

# Notes Updating Task
Analyze the conversation above to identify the user's preferences and how you can best satisfy them. Your goal is to create actionable notes that help you satisfy these preferences for future conversations. Keep your notes concise and actionable, without adding unnecessary details. Consider:
- When did the user explicitly ask you to adjust your response? What specifically did they want changed?
- What specific actions, formats, or approaches satisfy each preference? What should you keep in mind for future conversations?
As new situations arise, you may refine, combine, or split preferences to better reflect the user's needs. When updating the notes, do not lose any useful information from past interactions.
Make sure to add information about the user preferences that you are sure about, and do not hallucinate preferences.

# Output Format:
{{
   "user_preferences_reasoning": str, # Reasoning about the user preferences and how to satisfy them
   "agent_notes": str, # Updated notes. Provide a description of the user preferences, how to satisfy them, and any additional notes. This will be provided to you in future conversations with this user. Ensure that you provide a structured response that is clear and easy to understand.
}}
For each response, output a valid JSON object using the exact format above, do not include any text before or after the JSON object."""


def get_conversation_string(turns: List[Dict[str, str]]) -> str:
    """Convert conversation turns to string format."""
    conv_str = ""
    for turn in turns:
        role = "User" if turn["role"] == "user" else "Assistant"
        conv_str += f"{role}: {turn['content']}\n\n"
    return conv_str


def extract_enforced_preferences(full_user_log: List[Dict], conversation_turns: List[Dict]) -> List[str]:
    """
    Extract user messages where they enforce preferences.

    We identify enforcement by checking if:
    1. User's reasoning mentions enforcing/preference violation
    2. The response asks the agent to change behavior
    """
    enforced_messages = []

    # Get user turns from conversation
    user_turns = [t for t in conversation_turns if t["role"] == "user"]

    for i, log_entry in enumerate(full_user_log):
        reasoning = log_entry.get("reasoning", "").lower()

        # Indicators of preference enforcement
        enforcement_indicators = [
            "enforce",
            "preference",
            "not following",
            "should provide",
            "should use",
            "not adhering",
            "violat",
            "doesn't follow",
            "not address",
            "missing",
            "without",
            "need to",
        ]

        is_enforcement = any(ind in reasoning for ind in enforcement_indicators)

        # Also check if draft wasn't updated due to preference issues
        if "don't update" in reasoning.lower() or "do not update" in reasoning.lower():
            is_enforcement = True

        if is_enforcement and i < len(user_turns):
            enforced_messages.append(user_turns[i]["content"])

    return enforced_messages


def generate_reflection_target(conversation_turns: List[Dict], enforced_messages: List[str]) -> Dict:
    """
    Generate a target reflection based on the conversation.

    In practice, we'd use a strong LLM to generate gold reflections.
    For now, we create a structured template.
    """
    # Extract potential preferences from enforced messages
    preferences = []
    for msg in enforced_messages:
        # Simple heuristic extraction
        if "step" in msg.lower() or "process" in msg.lower():
            preferences.append({
                "preference": "Prefers step-by-step explanations",
                "context": "When solving problems",
                "action": "Break down solutions into clear numbered steps"
            })
        if "example" in msg.lower():
            preferences.append({
                "preference": "Prefers examples",
                "context": "When explaining concepts",
                "action": "Include concrete examples to illustrate points"
            })
        if "concise" in msg.lower() or "brief" in msg.lower():
            preferences.append({
                "preference": "Prefers concise responses",
                "context": "General communication",
                "action": "Keep responses focused and avoid unnecessary verbosity"
            })

    return {
        "user_preferences_reasoning": f"Based on the conversation, the user enforced {len(enforced_messages)} preferences.",
        "agent_notes": "\n".join([
            f"- {p['preference']}: {p['action']} ({p['context']})"
            for p in preferences
        ]) if preferences else "User did not express strong preferences in this conversation."
    }


def process_results_file(results_path: Path, output_dir: Path) -> Dict[str, int]:
    """Process a results.json file and generate training data."""

    with open(results_path) as f:
        data = json.load(f)

    sft_data = []
    grpo_data = []

    stats = {"total": 0, "with_enforcement": 0, "skipped": 0}

    for entry in data:
        stats["total"] += 1

        conversation = entry.get("conversation", {})
        turns = conversation.get("turns", [])
        full_user_log = entry.get("full_user_log", [])

        if len(turns) < 2:
            stats["skipped"] += 1
            continue

        # Extract enforced preferences
        enforced_messages = extract_enforced_preferences(full_user_log, turns)

        if enforced_messages:
            stats["with_enforcement"] += 1

        # Create conversation string
        conversation_str = get_conversation_string(turns)

        # Create input prompt
        input_prompt = update_agent_notes_prompt.format(
            agent_notes="No notes yet.",
            conversation_str=conversation_str
        )

        # Generate target reflection
        target_reflection = generate_reflection_target(turns, enforced_messages)

        # SFT format
        sft_entry = {
            "messages": [
                {"role": "user", "content": input_prompt},
                {"role": "assistant", "content": json.dumps(target_reflection, indent=2)}
            ]
        }
        sft_data.append(sft_entry)

        # GRPO format (includes enforced messages for reward)
        grpo_entry = {
            "messages": [
                {"role": "user", "content": input_prompt},
                {"role": "assistant", "content": json.dumps(target_reflection, indent=2)}
            ],
            "responses_that_enforce_preferences": enforced_messages,
            "profile_id": entry.get("profile_id", "unknown"),
            "problem_id": entry.get("problem_id", "unknown"),
        }
        grpo_data.append(grpo_entry)

    # Save data
    output_dir.mkdir(parents=True, exist_ok=True)

    sft_path = output_dir / "sft_training_data.json"
    with open(sft_path, "w") as f:
        json.dump(sft_data, f, indent=2)
    print(f"Saved {len(sft_data)} SFT examples to {sft_path}")

    grpo_path = output_dir / "grpo_training_data.json"
    with open(grpo_path, "w") as f:
        json.dump(grpo_data, f, indent=2)
    print(f"Saved {len(grpo_data)} GRPO examples to {grpo_path}")

    return stats


def main():
    parser = argparse.ArgumentParser(description="Generate training data from experiment results")
    parser.add_argument("--results-dir", type=str, required=True,
                       help="Directory containing results.json files")
    parser.add_argument("--output-dir", type=str, default="collaborativeagents/training/training_data",
                       help="Output directory for training data")
    args = parser.parse_args()

    results_dir = Path(args.results_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Find all results.json files
    results_files = list(results_dir.rglob("results.json"))
    print(f"Found {len(results_files)} results files")

    all_sft_data = []
    all_grpo_data = []
    total_stats = {"total": 0, "with_enforcement": 0, "skipped": 0}

    for results_file in results_files:
        print(f"\nProcessing: {results_file}")

        with open(results_file) as f:
            data = json.load(f)

        for entry in data:
            total_stats["total"] += 1

            conversation = entry.get("conversation", {})
            turns = conversation.get("turns", [])
            full_user_log = entry.get("full_user_log", [])

            if len(turns) < 2:
                total_stats["skipped"] += 1
                continue

            # Extract enforced preferences
            enforced_messages = extract_enforced_preferences(full_user_log, turns)

            if enforced_messages:
                total_stats["with_enforcement"] += 1

            # Create conversation string
            conversation_str = get_conversation_string(turns)

            # Create input prompt
            input_prompt = update_agent_notes_prompt.format(
                agent_notes="No notes yet.",
                conversation_str=conversation_str
            )

            # Generate target reflection
            target_reflection = generate_reflection_target(turns, enforced_messages)

            # SFT format
            sft_entry = {
                "messages": [
                    {"role": "user", "content": input_prompt},
                    {"role": "assistant", "content": json.dumps(target_reflection, indent=2)}
                ]
            }
            all_sft_data.append(sft_entry)

            # GRPO format (includes enforced messages for reward)
            grpo_entry = {
                "messages": [
                    {"role": "user", "content": input_prompt},
                    {"role": "assistant", "content": json.dumps(target_reflection, indent=2)}
                ],
                "responses_that_enforce_preferences": enforced_messages,
                "profile_id": entry.get("profile_id", "unknown"),
                "problem_id": entry.get("problem_id", "unknown"),
            }
            all_grpo_data.append(grpo_entry)

    # Save aggregated data
    sft_path = output_dir / "sft_training_data.json"
    with open(sft_path, "w") as f:
        json.dump(all_sft_data, f, indent=2)
    print(f"\nSaved {len(all_sft_data)} SFT examples to {sft_path}")

    grpo_path = output_dir / "grpo_training_data.json"
    with open(grpo_path, "w") as f:
        json.dump(all_grpo_data, f, indent=2)
    print(f"Saved {len(all_grpo_data)} GRPO examples to {grpo_path}")

    print(f"\n=== Summary ===")
    print(f"Total conversations: {total_stats['total']}")
    print(f"With enforcement: {total_stats['with_enforcement']}")
    print(f"Skipped: {total_stats['skipped']}")


if __name__ == "__main__":
    main()