src/multi_agent_datasets/multi_agent_loader.py

#!/usr/bin/env python3
"""
Multi-Agent Dataset Loader

This module provides comprehensive support for loading and processing multi-agent datasets
with two supported patterns:
A) Single folder with JSONLs that include an "agent" field
B) Per-agent subfolders (agent name == folder name)

Supports agent balancing, dataset validation, and integration with existing training pipelines.
"""

import os
import json
import yaml
import logging
from pathlib import Path
from typing import Dict, List, Optional, Union, Tuple, Any
from collections import Counter, defaultdict
from dataclasses import dataclass

import torch
from datasets import load_dataset, DatasetDict, Dataset, concatenate_datasets
from transformers import AutoTokenizer

logger = logging.getLogger(__name__)

@dataclass
class MultiAgentDatasetConfig:
    """Configuration for multi-agent dataset loading"""
    dataset_path: str
    agents_file: Optional[str] = None
    agent_prefix: str = "<|agent:"
    agent_suffix: str = "|>"
    balance_agents: bool = False
    balance_cap: Optional[int] = None
    max_seq_length: int = 2048
    validation_split: float = 0.1
    seed: int = 42

class MultiAgentDatasetLoader:
    """
    Multi-agent dataset loader supporting two patterns:
    1. Single folder with JSONLs containing 'agent' field
    2. Per-agent subfolders with agent name == folder name
    """
    
    def __init__(self, config: MultiAgentDatasetConfig):
        self.config = config
        self.agents = []
        self.dataset_stats = {}
        
    def read_agents_yaml(self, path: str) -> List[str]:
        """Read agents list from YAML file"""
        yml_path = os.path.join(path, "agents.yaml")
        if os.path.isfile(yml_path):
            try:
                with open(yml_path, "r") as f:
                    obj = yaml.safe_load(f) or {}
                    agents = [str(a) for a in obj.get("agents", [])]
                    logger.info(f"Loaded {len(agents)} agents from YAML: {agents}")
                    return agents
            except Exception as e:
                logger.warning(f"Failed to read agents.yaml: {e}")
        return []
    
    def list_agent_subdirs(self, path: str) -> List[Tuple[str, str]]:
        """List agent subdirectories with train/test.jsonl files"""
        items = []
        if not os.path.isdir(path):
            return items
            
        for name in sorted(os.listdir(path)):
            subdir_path = os.path.join(path, name)
            if os.path.isdir(subdir_path):
                train_file = os.path.join(subdir_path, "train.jsonl")
                test_file = os.path.join(subdir_path, "test.jsonl")
                if os.path.isfile(train_file) or os.path.isfile(test_file):
                    items.append((name, subdir_path))
                    logger.debug(f"Found agent subdirectory: {name}")
        
        return items
    
    def load_single_folder_dataset(self, dataset_path: str) -> DatasetDict:
        """Load dataset from single folder with agent field in rows"""
        data_files = {}
        
        train_file = os.path.join(dataset_path, "train.jsonl")
        test_file = os.path.join(dataset_path, "test.jsonl")
        
        if os.path.isfile(train_file):
            data_files["train"] = train_file
        if os.path.isfile(test_file):
            data_files["test"] = test_file
            
        if not data_files:
            raise FileNotFoundError(f"No dataset files found in {dataset_path}")
        
        logger.info(f"Loading single folder dataset from {data_files}")
        dataset = load_dataset("json", data_files=data_files)
        
        # Validate that agent field exists
        for split_name, split_data in dataset.items():
            if "agent" not in split_data.column_names:
                raise ValueError(f"Agent field not found in {split_name} split")
        
        return dataset
    
    def load_subfolder_dataset(self, dataset_path: str) -> DatasetDict:
        """Load dataset from per-agent subfolders"""
        subdirs = self.list_agent_subdirs(dataset_path)
        if not subdirs:
            raise FileNotFoundError(f"No agent subdirectories found in {dataset_path}")
        
        parts_train, parts_test = [], []
        
        for agent_name, agent_dir in subdirs:
            train_file = os.path.join(agent_dir, "train.jsonl")
            test_file = os.path.join(agent_dir, "test.jsonl")
            
            def add_agent_field(example):
                example["agent"] = agent_name
                return example
            
            if os.path.isfile(train_file):
                logger.debug(f"Loading train data for agent: {agent_name}")
                train_data = load_dataset("json", data_files={"train": train_file})["train"]
                train_data = train_data.map(add_agent_field)
                parts_train.append(train_data)
            
            if os.path.isfile(test_file):
                logger.debug(f"Loading test data for agent: {agent_name}")
                test_data = load_dataset("json", data_files={"test": test_file})["test"]
                test_data = test_data.map(add_agent_field)
                parts_test.append(test_data)
        
        dataset_dict = {}
        if parts_train:
            dataset_dict["train"] = concatenate_datasets(parts_train)
        if parts_test:
            dataset_dict["test"] = concatenate_datasets(parts_test)
        
        if not dataset_dict:
            raise ValueError("No data splits found in agent subdirectories")
        
        return DatasetDict(dataset_dict)
    
    def load_multiagent_dataset(self) -> DatasetDict:
        """
        Load multi-agent dataset supporting both patterns:
        - Single folder with 'agent' field in rows
        - Per-agent subfolders
        """
        dataset_path = self.config.dataset_path
        
        # Try single folder pattern first
        if os.path.isfile(os.path.join(dataset_path, "train.jsonl")):
            logger.info("Loading dataset using single folder pattern")
            return self.load_single_folder_dataset(dataset_path)
        
        # Try subfolder pattern
        logger.info("Loading dataset using subfolder pattern")
        return self.load_subfolder_dataset(dataset_path)
    
    def infer_agents_from_dataset(self, dataset: DatasetDict) -> List[str]:
        """Infer agent list from dataset"""
        agents = set()
        
        for split_name, split_data in dataset.items():
            if "agent" in split_data.column_names:
                agent_values = [a for a in set(split_data["agent"]) if a is not None]
                agents.update(agent_values)
                logger.debug(f"Found agents in {split_name}: {agent_values}")
        
        agents_list = sorted(list(agents))
        logger.info(f"Inferred {len(agents_list)} agents from dataset: {agents_list}")
        return agents_list
    
    def resolve_agents_list(self, dataset: DatasetDict) -> List[str]:
        """Resolve agents list from YAML file or dataset inference"""
        agents = []
        
        # Try to load from agents file first
        if self.config.agents_file and os.path.isfile(self.config.agents_file):
            try:
                with open(self.config.agents_file, "r") as f:
                    obj = yaml.safe_load(f) or {}
                    agents = [str(a) for a in obj.get("agents", [])]
                    logger.info(f"Loaded agents from file: {agents}")
            except Exception as e:
                logger.warning(f"Failed to load agents from file: {e}")
        
        # Fall back to dataset inference
        if not agents:
            agents = self.infer_agents_from_dataset(dataset)
        
        self.agents = agents
        return agents
    
    def balance_by_agent(self, dataset: Dataset, agent_col: str = "agent") -> Dataset:
        """
        Balance dataset by upsampling minority agents to the max count
        """
        if agent_col not in dataset.column_names:
            logger.warning(f"Agent column '{agent_col}' not found, skipping balancing")
            return dataset
        
        counts = Counter(dataset[agent_col])
        if not counts:
            logger.warning("No agent counts found, skipping balancing")
            return dataset
        
        max_count = max(counts.values())
        if self.config.balance_cap:
            max_count = min(max_count, self.config.balance_cap)
        
        logger.info(f"Balancing agents. Current counts: {dict(counts)}")
        logger.info(f"Target count per agent: {max_count}")
        
        parts = []
        for agent, count in counts.items():
            agent_subset = dataset.filter(lambda x: x[agent_col] == agent)
            parts.append(agent_subset)
            
            # Calculate how many additional samples we need
            needed = max_count - count
            if needed > 0:
                agent_subset_len = len(agent_subset)
                if agent_subset_len == 0:
                    logger.warning(f"Agent '{agent}' has zero samples, cannot upsample.")
                    continue
                # Calculate repetitions needed
                reps = needed // agent_subset_len
                remainder = needed % agent_subset_len
                
                # Add full repetitions
                for _ in range(reps):
                    parts.append(agent_subset)
                
                # Add remainder samples
                if remainder > 0:
                    remainder_subset = agent_subset.shuffle(seed=self.config.seed).select(range(remainder))
                    parts.append(remainder_subset)
        
        balanced_dataset = concatenate_datasets(parts).shuffle(seed=self.config.seed)
        
        # Log final counts
        final_counts = Counter(balanced_dataset[agent_col])
        logger.info(f"Balanced dataset counts: {dict(final_counts)}")
        
        return balanced_dataset
    
    def apply_agent_prefix(self, dataset: Dataset, tokenizer: AutoTokenizer) -> Dataset:
        """
        Apply agent prefix to dataset text using chat template or direct text
        """
        def add_agent_prefix(example):
            agent = example.get("agent", None)
            prefix = f"{self.config.agent_prefix}{agent}{self.config.agent_suffix}\n" if agent else ""
            
            # Handle different text formats
            if "messages" in example and example["messages"] is not None:
                # Use chat template if available
                try:
                    text = tokenizer.apply_chat_template(
                        example["messages"], 
                        tokenize=False, 
                        add_generation_prompt=False
                    )
                    example["text"] = prefix + text
                except Exception as e:
                    logger.warning(f"Failed to apply chat template: {e}")
                    # Fallback to simple concatenation
                    text = str(example["messages"])
                    example["text"] = prefix + text
                    
            elif "text" in example and example["text"] is not None:
                example["text"] = prefix + example["text"]
                
            else:
                # Handle prompt/response format
                prompt = example.get("prompt", "")
                response = example.get("response", "")
                example["text"] = prefix + prompt + ("\n" if response else "") + response
            
            return example
        
        # Get original features to preserve
        original_features = list(dataset.features)
        features_to_remove = [f for f in original_features if f not in ["text", "agent"]]
        
        logger.info("Applying agent prefixes to dataset")
        processed_dataset = dataset.map(
            add_agent_prefix,
            remove_columns=features_to_remove,
            desc="Adding agent prefixes"
        )
        
        return processed_dataset
    
    def validate_dataset(self, dataset: DatasetDict) -> Dict[str, Any]:
        """Validate dataset and return statistics"""
        stats = {
            "total_samples": 0,
            "agents": {},
            "splits": {},
            "validation_errors": []
        }
        
        for split_name, split_data in dataset.items():
            split_stats = {
                "samples": len(split_data),
                "agents": {},
                "columns": split_data.column_names
            }
            
            stats["total_samples"] += len(split_data)
            
            # Validate required columns
            if "agent" not in split_data.column_names:
                stats["validation_errors"].append(f"Missing 'agent' column in {split_name}")
            
            # Count agents
            if "agent" in split_data.column_names:
                agent_counts = Counter(split_data["agent"])
                split_stats["agents"] = dict(agent_counts)
                
                # Update global agent counts
                for agent, count in agent_counts.items():
                    if agent not in stats["agents"]:
                        stats["agents"][agent] = 0
                    stats["agents"][agent] += count
            
            stats["splits"][split_name] = split_stats
        
        self.dataset_stats = stats
        logger.info(f"Dataset validation complete. Stats: {stats}")
        
        return stats
    
    def load_and_process(self, tokenizer: AutoTokenizer) -> Tuple[DatasetDict, List[str], Dict[str, Any]]:
        """
        Complete dataset loading and processing pipeline
        """
        logger.info(f"Loading multi-agent dataset from {self.config.dataset_path}")
        
        # Load dataset
        dataset = self.load_multiagent_dataset()
        
        # Resolve agents list
        agents = self.resolve_agents_list(dataset)
        
        # Validate dataset
        stats = self.validate_dataset(dataset)
        
        # Apply agent prefixes
        if "train" in dataset:
            dataset["train"] = self.apply_agent_prefix(dataset["train"], tokenizer)
        if "test" in dataset:
            dataset["test"] = self.apply_agent_prefix(dataset["test"], tokenizer)
        
        # Balance agents if requested
        if self.config.balance_agents and "train" in dataset:
            dataset["train"] = self.balance_by_agent(dataset["train"])
        
        logger.info(f"Dataset processing complete. Loaded {len(agents)} agents with {stats['total_samples']} total samples")
        
        return dataset, agents, stats

class MultiAgentDatasetValidator:
    """Validator for multi-agent datasets"""
    
    @staticmethod
    def validate_jsonl_file(file_path: str) -> List[str]:
        """Validate JSONL file format and content"""
        errors = []
        
        if not os.path.isfile(file_path):
            errors.append(f"File not found: {file_path}")
            return errors
        
        try:
            with open(file_path, 'r') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    
                    try:
                        data = json.loads(line)
                        
                        # Check required fields
                        if not isinstance(data, dict):
                            errors.append(f"Line {line_num}: Not a JSON object")
                            continue
                        
                        # Check for agent field
                        if "agent" not in data:
                            errors.append(f"Line {line_num}: Missing 'agent' field")
                        
                        # Check for text content
                        has_text = any(field in data for field in ["text", "messages", "prompt"])
                        if not has_text:
                            errors.append(f"Line {line_num}: No text content found")
                        
                    except json.JSONDecodeError as e:
                        errors.append(f"Line {line_num}: JSON decode error - {e}")
        
        except Exception as e:
            errors.append(f"File read error: {e}")
        
        return errors
    
    @staticmethod
    def validate_dataset_structure(dataset_path: str) -> Dict[str, Any]:
        """Validate complete dataset structure"""
        validation_result = {
            "valid": True,
            "errors": [],
            "warnings": [],
            "structure": {}
        }
        
        if not os.path.isdir(dataset_path):
            validation_result["valid"] = False
            validation_result["errors"].append(f"Dataset path is not a directory: {dataset_path}")
            return validation_result
        
        # Check for single folder pattern
        train_file = os.path.join(dataset_path, "train.jsonl")
        test_file = os.path.join(dataset_path, "test.jsonl")
        
        if os.path.isfile(train_file):
            validation_result["structure"]["pattern"] = "single_folder"
            validation_result["structure"]["files"] = []
            
            if os.path.isfile(train_file):
                validation_result["structure"]["files"].append("train.jsonl")
                errors = MultiAgentDatasetValidator.validate_jsonl_file(train_file)
                validation_result["errors"].extend(errors)
            
            if os.path.isfile(test_file):
                validation_result["structure"]["files"].append("test.jsonl")
                errors = MultiAgentDatasetValidator.validate_jsonl_file(test_file)
                validation_result["errors"].extend(errors)
        
        else:
            # Check for subfolder pattern
            validation_result["structure"]["pattern"] = "subfolders"
            validation_result["structure"]["agents"] = []
            
            for item in os.listdir(dataset_path):
                item_path = os.path.join(dataset_path, item)
                if os.path.isdir(item_path):
                    agent_train = os.path.join(item_path, "train.jsonl")
                    agent_test = os.path.join(item_path, "test.jsonl")
                    
                    if os.path.isfile(agent_train) or os.path.isfile(agent_test):
                        validation_result["structure"]["agents"].append(item)
                        
                        if os.path.isfile(agent_train):
                            errors = MultiAgentDatasetValidator.validate_jsonl_file(agent_train)
                            validation_result["errors"].extend([f"{item}/train.jsonl: {e}" for e in errors])
                        
                        if os.path.isfile(agent_test):
                            errors = MultiAgentDatasetValidator.validate_jsonl_file(agent_test)
                            validation_result["errors"].extend([f"{item}/test.jsonl: {e}" for e in errors])
        
        # Check for agents.yaml
        agents_yaml = os.path.join(dataset_path, "agents.yaml")
        if os.path.isfile(agents_yaml):
            validation_result["structure"]["has_agents_yaml"] = True
            try:
                with open(agents_yaml, 'r') as f:
                    yaml.safe_load(f)
            except Exception as e:
                validation_result["warnings"].append(f"Invalid agents.yaml: {e}")
        else:
            validation_result["structure"]["has_agents_yaml"] = False
        
        validation_result["valid"] = len(validation_result["errors"]) == 0
        
        return validation_result

# Example usage and testing
if __name__ == "__main__":
    # Configure logging
    logging.basicConfig(level=logging.INFO)
    
    # Example configuration
    config = MultiAgentDatasetConfig(
        dataset_path="/path/to/dataset",
        balance_agents=True,
        balance_cap=1000
    )
    
    # Create loader
    loader = MultiAgentDatasetLoader(config)
    
    # Example tokenizer (would be loaded from actual model)
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-MoE-instruct")
    
    try:
        # Load and process dataset
        dataset, agents, stats = loader.load_and_process(tokenizer)
        
        print(f"Loaded dataset with {len(agents)} agents:")
        for agent in agents:
            print(f"  - {agent}")
        
        print(f"Dataset stats: {stats}")
        
    except Exception as e:
        print(f"Error loading dataset: {e}")