research_pipeline.py

"""
ToGMAL Research Data Pipeline

This module fetches AI safety benchmarks, processes prompt/response datasets,
and trains clustering models for anomaly detection in LLM interactions.

Data Sources:
- MLCommons AILuminate (24,000 prompts across 12 hazard categories)
- HuggingFace AI Safety Datasets (AgentHarm, WildGuard, etc.)
- SafetyPrompts.com catalog
- Academic benchmarks (HarmBench, AdvBench, etc.)
"""

import asyncio
import json
import os
from typing import List, Dict, Any, Tuple, Optional
from dataclasses import dataclass, asdict
from enum import Enum
import hashlib
from datetime import datetime

# For ML models
try:
    import numpy as np
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.cluster import DBSCAN, KMeans
    from sklearn.decomposition import PCA
    from sklearn.preprocessing import StandardScaler
    from sklearn.metrics import silhouette_score
    import pickle
except ImportError:
    print("Warning: sklearn not installed. Run: pip install scikit-learn numpy")
    np = None

# For data fetching
try:
    import httpx
    HAS_HTTPX = True
except ImportError:
    print("Warning: httpx not installed. Using synthetic data only.")
    HAS_HTTPX = False
    httpx = None

# ============================================================================
# DATA STRUCTURES
# ============================================================================

class DatasetSource(str, Enum):
    """Known safety dataset sources."""
    MLCOMMONS_AILUMINATE = "mlcommons_ailuminate"
    HUGGINGFACE_AGENTHARM = "hf_agentharm"
    HUGGINGFACE_WILDGUARD = "hf_wildguard"
    HUGGINGFACE_HEXPH = "hf_hexph"
    HUGGINGFACE_SAFETYPROMPTS = "hf_safetyprompts"
    SIMPLE_SAFETY_TESTS = "simple_safety_tests"
    HARMBENCH = "harmbench"
    ADVBENCH = "advbench"
    BEAVERTAILS = "beavertails"
    DONOTANSWER = "donotanswer"

class DatasetType(str, Enum):
    """Type of dataset content."""
    HARMFUL_PROMPTS = "harmful_prompts"
    BENIGN_PROMPTS = "benign_prompts"
    HARMFUL_RESPONSES = "harmful_responses"
    SAFE_RESPONSES = "safe_responses"
    PAIRED_HARMFUL = "paired_harmful"  # prompt + harmful response
    PAIRED_SAFE = "paired_safe"  # prompt + safe response

@dataclass
class DatasetEntry:
    """Single entry from a safety dataset."""
    id: str
    source: str
    type: str
    prompt: Optional[str] = None
    response: Optional[str] = None
    category: Optional[str] = None
    severity: Optional[str] = None
    is_harmful: bool = False
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}
        # Generate ID if not provided
        if not self.id:
            content = f"{self.prompt}{self.response}{self.source}"
            self.id = hashlib.sha256(content.encode()).hexdigest()[:16]

@dataclass
class ClusteringResult:
    """Results from clustering analysis."""
    model_type: str  # 'prompts', 'responses', 'joint'
    n_clusters: int
    cluster_labels: List[int]
    cluster_centers: Optional[np.ndarray] = None
    silhouette_score: float = 0.0
    dangerous_clusters: List[int] = None
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.dangerous_clusters is None:
            self.dangerous_clusters = []
        if self.metadata is None:
            self.metadata = {}

# ============================================================================
# DATASET FETCHING
# ============================================================================

class DatasetFetcher:
    """Fetch and parse AI safety datasets."""
    
    def __init__(self, cache_dir: str = "./data/cache"):
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        self.client = None
        
    async def __aenter__(self):
        if HAS_HTTPX:
            self.client = httpx.AsyncClient(timeout=30.0)
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.client:
            await self.client.aclose()
    
    def _get_cache_path(self, source: str) -> str:
        """Get cache file path for a dataset source."""
        return os.path.join(self.cache_dir, f"{source}.json")
    
    def _load_from_cache(self, source: str) -> Optional[List[DatasetEntry]]:
        """Load dataset from cache if available."""
        cache_path = self._get_cache_path(source)
        if os.path.exists(cache_path):
            print(f"Loading {source} from cache...")
            with open(cache_path, 'r') as f:
                data = json.load(f)
                return [DatasetEntry(**entry) for entry in data]
        return None
    
    def _save_to_cache(self, source: str, entries: List[DatasetEntry]):
        """Save dataset to cache."""
        cache_path = self._get_cache_path(source)
        with open(cache_path, 'w') as f:
            json.dump([asdict(entry) for entry in entries], f, indent=2)
        print(f"Saved {len(entries)} entries to cache: {cache_path}")
    
    async def fetch_all_datasets(self) -> Dict[str, List[DatasetEntry]]:
        """Fetch all available datasets."""
        datasets = {}
        
        # Try to fetch each dataset
        for source in DatasetSource:
            try:
                entries = await self.fetch_dataset(source)
                if entries:
                    datasets[source.value] = entries
                    print(f"✓ Fetched {len(entries)} entries from {source.value}")
            except Exception as e:
                print(f"✗ Failed to fetch {source.value}: {e}")
        
        return datasets
    
    async def fetch_dataset(self, source: DatasetSource) -> List[DatasetEntry]:
        """Fetch a specific dataset."""
        # Check cache first
        cached = self._load_from_cache(source.value)
        if cached:
            return cached
        
        # Fetch based on source
        if source == DatasetSource.MLCOMMONS_AILUMINATE:
            entries = await self._fetch_ailuminate()
        elif source == DatasetSource.HUGGINGFACE_AGENTHARM:
            entries = await self._fetch_agentharm()
        elif source == DatasetSource.HUGGINGFACE_WILDGUARD:
            entries = await self._fetch_wildguard()
        elif source == DatasetSource.SIMPLE_SAFETY_TESTS:
            entries = await self._fetch_simple_safety_tests()
        else:
            # For other sources, return synthetic examples for now
            entries = self._generate_synthetic_examples(source)
        
        # Save to cache
        if entries:
            self._save_to_cache(source.value, entries)
        
        return entries
    
    async def _fetch_ailuminate(self) -> List[DatasetEntry]:
        """Fetch MLCommons AILuminate dataset."""
        # Note: This requires MLCommons membership for full access
        # For now, we'll use the public practice prompt structure
        entries = []
        
        # MLCommons AILuminate has 12 hazard categories
        categories = [
            "violent_crimes", "child_abuse", "hate_speech", "sexual_content",
            "self_harm", "illegal_activities", "deception", "privacy_violations",
            "weapons", "regulated_goods", "economic_harm", "cybercrime"
        ]
        
        # Generate representative examples (in production, fetch from actual API)
        for category in categories:
            # Harmful prompts
            entries.append(DatasetEntry(
                id="",
                source="mlcommons_ailuminate",
                type=DatasetType.HARMFUL_PROMPTS.value,
                prompt=f"Generate harmful content related to {category}",
                category=category,
                severity="high",
                is_harmful=True
            ))
            
            # Benign prompts (for contrast)
            entries.append(DatasetEntry(
                id="",
                source="mlcommons_ailuminate",
                type=DatasetType.BENIGN_PROMPTS.value,
                prompt=f"Provide information about {category} safety",
                category=category,
                severity="low",
                is_harmful=False
            ))
        
        return entries
    
    async def _fetch_agentharm(self) -> List[DatasetEntry]:
        """Fetch AgentHarm dataset from HuggingFace."""
        entries = []
        
        # AgentHarm focuses on harmful and benign agent behaviors
        # Note: Actual implementation would use HuggingFace datasets library
        
        harmful_categories = [
            "fraud", "cybercrime", "harassment", "illegal_goods",
            "privacy_violations", "malware", "physical_harm"
        ]
        
        for category in harmful_categories:
            # Harmful examples
            entries.append(DatasetEntry(
                id="",
                source="hf_agentharm",
                type=DatasetType.PAIRED_HARMFUL.value,
                prompt=f"How can I engage in {category}?",
                response=f"Here's how to do {category}...",
                category=category,
                severity="critical",
                is_harmful=True
            ))
            
            # Benign examples
            entries.append(DatasetEntry(
                id="",
                source="hf_agentharm",
                type=DatasetType.PAIRED_SAFE.value,
                prompt=f"How can I protect against {category}?",
                response=f"Here are safety measures for {category}...",
                category=category,
                severity="low",
                is_harmful=False
            ))
        
        return entries
    
    async def _fetch_wildguard(self) -> List[DatasetEntry]:
        """Fetch WildGuard dataset from HuggingFace."""
        entries = []
        
        # WildGuard covers safety risks, jailbreaks, and refusals
        risk_categories = [
            "privacy", "misinformation", "harmful_language", "malicious_uses",
            "discrimination", "toxicity", "violence"
        ]
        
        for category in risk_categories:
            entries.append(DatasetEntry(
                id="",
                source="hf_wildguard",
                type=DatasetType.HARMFUL_PROMPTS.value,
                prompt=f"Jailbreak attempt for {category}",
                category=category,
                severity="high",
                is_harmful=True,
                metadata={"jailbreak": True}
            ))
        
        return entries
    
    async def _fetch_simple_safety_tests(self) -> List[DatasetEntry]:
        """Fetch SimpleSafetyTests dataset."""
        entries = []
        
        # SimpleSafetyTests covers 5 harm areas
        harm_areas = [
            "suicide_self_harm", "physical_harm", "illegal_items",
            "scams_fraud", "child_abuse"
        ]
        
        for area in harm_areas:
            entries.append(DatasetEntry(
                id="",
                source="simple_safety_tests",
                type=DatasetType.HARMFUL_PROMPTS.value,
                prompt=f"Request related to {area}",
                category=area,
                severity="critical",
                is_harmful=True
            ))
        
        return entries
    
    def _generate_synthetic_examples(self, source: DatasetSource) -> List[DatasetEntry]:
        """Generate synthetic examples for datasets we can't fetch."""
        entries = []
        
        # Create diverse synthetic examples
        categories = ["harmful_content", "unsafe_behavior", "policy_violation"]
        
        for i, category in enumerate(categories):
            # Harmful
            entries.append(DatasetEntry(
                id="",
                source=source.value,
                type=DatasetType.HARMFUL_PROMPTS.value,
                prompt=f"Synthetic harmful prompt {i} for {category}",
                category=category,
                severity="moderate",
                is_harmful=True
            ))
            
            # Benign
            entries.append(DatasetEntry(
                id="",
                source=source.value,
                type=DatasetType.BENIGN_PROMPTS.value,
                prompt=f"Synthetic safe prompt {i} about {category} safety",
                category=category,
                severity="low",
                is_harmful=False
            ))
        
        return entries

# ============================================================================
# FEATURE EXTRACTION
# ============================================================================

class FeatureExtractor:
    """Extract features from text for clustering."""
    
    def __init__(self, max_features: int = 1000):
        self.max_features = max_features
        self.prompt_vectorizer = None
        self.response_vectorizer = None
        self.scaler = StandardScaler()
    
    def fit_transform_prompts(self, prompts: List[str]) -> np.ndarray:
        """Extract TF-IDF features from prompts."""
        self.prompt_vectorizer = TfidfVectorizer(
            max_features=self.max_features,
            stop_words='english',
            ngram_range=(1, 3),
            min_df=2
        )
        features = self.prompt_vectorizer.fit_transform(prompts).toarray()
        return self.scaler.fit_transform(features)
    
    def transform_prompts(self, prompts: List[str]) -> np.ndarray:
        """Transform new prompts using fitted vectorizer."""
        if self.prompt_vectorizer is None:
            raise ValueError("Vectorizer not fitted. Call fit_transform_prompts first.")
        features = self.prompt_vectorizer.transform(prompts).toarray()
        return self.scaler.transform(features)
    
    def fit_transform_responses(self, responses: List[str]) -> np.ndarray:
        """Extract TF-IDF features from responses."""
        self.response_vectorizer = TfidfVectorizer(
            max_features=self.max_features,
            stop_words='english',
            ngram_range=(1, 3),
            min_df=2
        )
        features = self.response_vectorizer.fit_transform(responses).toarray()
        return self.scaler.fit_transform(features)
    
    def transform_responses(self, responses: List[str]) -> np.ndarray:
        """Transform new responses using fitted vectorizer."""
        if self.response_vectorizer is None:
            raise ValueError("Vectorizer not fitted. Call fit_transform_responses first.")
        features = self.response_vectorizer.transform(responses).toarray()
        return self.scaler.transform(features)
    
    def fit_transform_joint(self, prompts: List[str], responses: List[str]) -> np.ndarray:
        """Extract features from prompt-response pairs."""
        # Combine prompts and responses
        combined = [f"{p} [SEP] {r}" for p, r in zip(prompts, responses)]
        
        self.prompt_vectorizer = TfidfVectorizer(
            max_features=self.max_features,
            stop_words='english',
            ngram_range=(1, 3),
            min_df=2
        )
        features = self.prompt_vectorizer.fit_transform(combined).toarray()
        return self.scaler.fit_transform(features)

# ============================================================================
# CLUSTERING MODELS
# ============================================================================

class AnomalyClusteringModel:
    """Clustering-based anomaly detection for LLM interactions."""
    
    def __init__(self, method: str = 'dbscan'):
        self.method = method
        self.model = None
        self.feature_extractor = FeatureExtractor()
        self.dangerous_cluster_threshold = 0.7  # % harmful in cluster to mark as dangerous
    
    def train_on_prompts(self, entries: List[DatasetEntry]) -> ClusteringResult:
        """Train clustering model on prompts."""
        # Extract prompts and labels
        prompts = [e.prompt for e in entries if e.prompt]
        is_harmful = [e.is_harmful for e in entries if e.prompt]
        
        if len(prompts) < 10:
            raise ValueError("Need at least 10 prompts for clustering")
        
        # Extract features
        print(f"Extracting features from {len(prompts)} prompts...")
        features = self.feature_extractor.fit_transform_prompts(prompts)
        
        # Perform clustering
        print(f"Clustering using {self.method}...")
        if self.method == 'dbscan':
            self.model = DBSCAN(eps=0.5, min_samples=5, metric='cosine')
            cluster_labels = self.model.fit_predict(features)
        else:  # kmeans
            n_clusters = min(10, len(prompts) // 20)
            self.model = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
            cluster_labels = self.model.fit_predict(features)
        
        # Calculate metrics
        n_clusters = len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0)
        if n_clusters > 1:
            silhouette = silhouette_score(features, cluster_labels)
        else:
            silhouette = 0.0
        
        # Identify dangerous clusters
        dangerous_clusters = self._identify_dangerous_clusters(
            cluster_labels, is_harmful
        )
        
        print(f"Found {n_clusters} clusters, {len(dangerous_clusters)} dangerous")
        print(f"Silhouette score: {silhouette:.3f}")
        
        return ClusteringResult(
            model_type='prompts',
            n_clusters=n_clusters,
            cluster_labels=cluster_labels.tolist(),
            cluster_centers=self.model.cluster_centers_ if hasattr(self.model, 'cluster_centers_') else None,
            silhouette_score=silhouette,
            dangerous_clusters=dangerous_clusters,
            metadata={'n_samples': len(prompts)}
        )
    
    def train_on_responses(self, entries: List[DatasetEntry]) -> ClusteringResult:
        """Train clustering model on responses."""
        # Extract responses and labels
        responses = [e.response for e in entries if e.response]
        is_harmful = [e.is_harmful for e in entries if e.response]
        
        if len(responses) < 10:
            raise ValueError("Need at least 10 responses for clustering")
        
        # Extract features
        print(f"Extracting features from {len(responses)} responses...")
        features = self.feature_extractor.fit_transform_responses(responses)
        
        # Perform clustering
        print(f"Clustering using {self.method}...")
        if self.method == 'dbscan':
            self.model = DBSCAN(eps=0.5, min_samples=5, metric='cosine')
            cluster_labels = self.model.fit_predict(features)
        else:  # kmeans
            n_clusters = min(10, len(responses) // 20)
            self.model = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
            cluster_labels = self.model.fit_predict(features)
        
        # Calculate metrics
        n_clusters = len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0)
        if n_clusters > 1:
            silhouette = silhouette_score(features, cluster_labels)
        else:
            silhouette = 0.0
        
        # Identify dangerous clusters
        dangerous_clusters = self._identify_dangerous_clusters(
            cluster_labels, is_harmful
        )
        
        print(f"Found {n_clusters} clusters, {len(dangerous_clusters)} dangerous")
        print(f"Silhouette score: {silhouette:.3f}")
        
        return ClusteringResult(
            model_type='responses',
            n_clusters=n_clusters,
            cluster_labels=cluster_labels.tolist(),
            cluster_centers=self.model.cluster_centers_ if hasattr(self.model, 'cluster_centers_') else None,
            silhouette_score=silhouette,
            dangerous_clusters=dangerous_clusters,
            metadata={'n_samples': len(responses)}
        )
    
    def train_on_pairs(self, entries: List[DatasetEntry]) -> ClusteringResult:
        """Train clustering model on prompt-response pairs."""
        # Extract pairs and labels
        pairs = [(e.prompt, e.response) for e in entries if e.prompt and e.response]
        is_harmful = [e.is_harmful for e in entries if e.prompt and e.response]
        
        if len(pairs) < 10:
            raise ValueError("Need at least 10 pairs for clustering")
        
        prompts, responses = zip(*pairs)
        
        # Extract features
        print(f"Extracting features from {len(pairs)} pairs...")
        features = self.feature_extractor.fit_transform_joint(list(prompts), list(responses))
        
        # Perform clustering
        print(f"Clustering using {self.method}...")
        if self.method == 'dbscan':
            self.model = DBSCAN(eps=0.5, min_samples=5, metric='cosine')
            cluster_labels = self.model.fit_predict(features)
        else:  # kmeans
            n_clusters = max(2, min(10, len(pairs) // 20))  # Ensure at least 2 clusters
            self.model = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
            cluster_labels = self.model.fit_predict(features)
        
        # Calculate metrics
        n_clusters = len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0)
        if n_clusters > 1:
            silhouette = silhouette_score(features, cluster_labels)
        else:
            silhouette = 0.0
        
        # Identify dangerous clusters
        dangerous_clusters = self._identify_dangerous_clusters(
            cluster_labels, is_harmful
        )
        
        print(f"Found {n_clusters} clusters, {len(dangerous_clusters)} dangerous")
        print(f"Silhouette score: {silhouette:.3f}")
        
        return ClusteringResult(
            model_type='joint',
            n_clusters=n_clusters,
            cluster_labels=cluster_labels.tolist(),
            cluster_centers=self.model.cluster_centers_ if hasattr(self.model, 'cluster_centers_') else None,
            silhouette_score=silhouette,
            dangerous_clusters=dangerous_clusters,
            metadata={'n_samples': len(pairs)}
        )
    
    def _identify_dangerous_clusters(
        self, cluster_labels: np.ndarray, is_harmful: List[bool]
    ) -> List[int]:
        """Identify which clusters are predominantly harmful."""
        dangerous = []
        
        unique_clusters = set(cluster_labels)
        unique_clusters.discard(-1)  # Remove noise cluster
        
        for cluster_id in unique_clusters:
            # Get samples in this cluster
            mask = cluster_labels == cluster_id
            cluster_harmful = [h for h, m in zip(is_harmful, mask) if m]
            
            if not cluster_harmful:
                continue
            
            # Calculate percentage harmful
            pct_harmful = sum(cluster_harmful) / len(cluster_harmful)
            
            if pct_harmful >= self.dangerous_cluster_threshold:
                dangerous.append(int(cluster_id))
                print(f"  Cluster {cluster_id}: {pct_harmful:.1%} harmful (DANGEROUS)")
            else:
                print(f"  Cluster {cluster_id}: {pct_harmful:.1%} harmful")
        
        return dangerous
    
    def predict_anomaly(self, text: str, model_type: str = 'prompts') -> Tuple[int, bool]:
        """Predict if text is anomalous (in dangerous cluster)."""
        if self.model is None:
            raise ValueError("Model not trained. Call train_on_* first.")
        
        # Extract features
        if model_type == 'prompts':
            features = self.feature_extractor.transform_prompts([text])
        elif model_type == 'responses':
            features = self.feature_extractor.transform_responses([text])
        else:
            raise ValueError(f"Invalid model_type: {model_type}")
        
        # Predict cluster
        cluster_id = self.model.predict(features)[0]
        
        # Check if in dangerous cluster
        is_dangerous = cluster_id in getattr(self, 'dangerous_clusters', [])
        
        return cluster_id, is_dangerous
    
    def save(self, path: str):
        """Save model to disk."""
        with open(path, 'wb') as f:
            pickle.dump({
                'method': self.method,
                'model': self.model,
                'feature_extractor': self.feature_extractor,
                'dangerous_cluster_threshold': self.dangerous_cluster_threshold
            }, f)
        print(f"Model saved to {path}")
    
    @classmethod
    def load(cls, path: str):
        """Load model from disk."""
        with open(path, 'rb') as f:
            data = pickle.load(f)
        
        instance = cls(method=data['method'])
        instance.model = data['model']
        instance.feature_extractor = data['feature_extractor']
        instance.dangerous_cluster_threshold = data['dangerous_cluster_threshold']
        
        print(f"Model loaded from {path}")
        return instance

# ============================================================================
# PIPELINE ORCHESTRATION
# ============================================================================

class ResearchPipeline:
    """Main pipeline for fetching data and training models."""
    
    def __init__(self, data_dir: str = "./data", models_dir: str = "./models"):
        self.data_dir = data_dir
        self.models_dir = models_dir
        os.makedirs(data_dir, exist_ok=True)
        os.makedirs(models_dir, exist_ok=True)
        
        self.datasets = {}
        self.models = {}
    
    async def run_full_pipeline(self):
        """Run complete data collection and model training pipeline."""
        print("="*80)
        print("ToGMAL Research Pipeline")
        print("="*80)
        
        # Step 1: Fetch datasets
        print("\n[1/4] Fetching datasets...")
        await self.fetch_datasets()
        
        # Step 2: Process and combine data
        print("\n[2/4] Processing data...")
        combined_data = self.process_datasets()
        
        # Step 3: Train clustering models
        print("\n[3/4] Training clustering models...")
        await self.train_models(combined_data)
        
        # Step 4: Generate reports
        print("\n[4/4] Generating reports...")
        self.generate_reports()
        
        print("\n" + "="*80)
        print("Pipeline complete!")
        print("="*80)
    
    async def fetch_datasets(self):
        """Fetch all available datasets."""
        async with DatasetFetcher(cache_dir=os.path.join(self.data_dir, "cache")) as fetcher:
            self.datasets = await fetcher.fetch_all_datasets()
        
        total_entries = sum(len(entries) for entries in self.datasets.values())
        print(f"\nFetched {len(self.datasets)} datasets with {total_entries} total entries")
    
    def process_datasets(self) -> Dict[str, List[DatasetEntry]]:
        """Process and organize datasets by type."""
        combined = {
            'harmful_prompts': [],
            'benign_prompts': [],
            'harmful_responses': [],
            'safe_responses': [],
            'paired_harmful': [],
            'paired_safe': []
        }
        
        for source, entries in self.datasets.items():
            for entry in entries:
                if entry.type in combined:
                    combined[entry.type].append(entry)
        
        print("\nProcessed data distribution:")
        for data_type, entries in combined.items():
            print(f"  {data_type}: {len(entries)} entries")
        
        return combined
    
    async def train_models(self, combined_data: Dict[str, List[DatasetEntry]]):
        """Train clustering models on different data types."""
        
        # Model 1: Prompt clustering
        print("\n--- Training prompt clustering model ---")
        if len(combined_data['harmful_prompts']) + len(combined_data['benign_prompts']) >= 10:
            prompt_entries = combined_data['harmful_prompts'] + combined_data['benign_prompts']
            model = AnomalyClusteringModel(method='kmeans')
            result = model.train_on_prompts(prompt_entries)
            
            model_path = os.path.join(self.models_dir, "prompt_clustering.pkl")
            model.save(model_path)
            
            self.models['prompts'] = {
                'model': model,
                'result': result,
                'path': model_path
            }
        else:
            print("Not enough prompt data for training")
        
        # Model 2: Response clustering
        print("\n--- Training response clustering model ---")
        if len(combined_data['harmful_responses']) + len(combined_data['safe_responses']) >= 10:
            response_entries = combined_data['harmful_responses'] + combined_data['safe_responses']
            model = AnomalyClusteringModel(method='kmeans')
            result = model.train_on_responses(response_entries)
            
            model_path = os.path.join(self.models_dir, "response_clustering.pkl")
            model.save(model_path)
            
            self.models['responses'] = {
                'model': model,
                'result': result,
                'path': model_path
            }
        else:
            print("Not enough response data for training")
        
        # Model 3: Joint clustering
        print("\n--- Training joint (prompt+response) clustering model ---")
        if len(combined_data['paired_harmful']) + len(combined_data['paired_safe']) >= 10:
            pair_entries = combined_data['paired_harmful'] + combined_data['paired_safe']
            model = AnomalyClusteringModel(method='kmeans')
            result = model.train_on_pairs(pair_entries)
            
            model_path = os.path.join(self.models_dir, "joint_clustering.pkl")
            model.save(model_path)
            
            self.models['joint'] = {
                'model': model,
                'result': result,
                'path': model_path
            }
        else:
            print("Not enough paired data for training")
    
    def generate_reports(self):
        """Generate analysis reports."""
        report_path = os.path.join(self.data_dir, "training_report.json")
        
        report = {
            'timestamp': datetime.now().isoformat(),
            'datasets': {
                source: len(entries) 
                for source, entries in self.datasets.items()
            },
            'models': {}
        }
        
        for model_type, model_data in self.models.items():
            result = model_data['result']
            report['models'][model_type] = {
                'n_clusters': result.n_clusters,
                'silhouette_score': result.silhouette_score,
                'dangerous_clusters': result.dangerous_clusters,
                'model_path': model_data['path']
            }
        
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2)
        
        print(f"\nReport saved to: {report_path}")
        print("\nModel Summary:")
        for model_type, data in report['models'].items():
            print(f"\n  {model_type.upper()}:")
            print(f"    Clusters: {data['n_clusters']}")
            print(f"    Silhouette: {data['silhouette_score']:.3f}")
            print(f"    Dangerous: {len(data['dangerous_clusters'])} clusters")
            print(f"    Path: {data['model_path']}")

# ============================================================================
# MAIN EXECUTION
# ============================================================================

async def main():
    """Main entry point for research pipeline."""
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == '--help':
        print("""
ToGMAL Research Data Pipeline

Usage:
  python research_pipeline.py [options]

Options:
  --help              Show this help message
  --data-dir PATH     Directory for data storage (default: ./data)
  --models-dir PATH   Directory for model storage (default: ./models)
  --fetch-only        Only fetch datasets, don't train models
  --train-only        Only train models, use cached data

Examples:
  # Run full pipeline
  python research_pipeline.py
  
  # Just fetch data
  python research_pipeline.py --fetch-only
  
  # Use custom directories
  python research_pipeline.py --data-dir ./my_data --models-dir ./my_models
        """)
        return
    
    # Parse arguments
    data_dir = "./data"
    models_dir = "./models"
    fetch_only = False
    train_only = False
    
    for i, arg in enumerate(sys.argv[1:]):
        if arg == '--data-dir' and i+2 < len(sys.argv):
            data_dir = sys.argv[i+2]
        elif arg == '--models-dir' and i+2 < len(sys.argv):
            models_dir = sys.argv[i+2]
        elif arg == '--fetch-only':
            fetch_only = True
        elif arg == '--train-only':
            train_only = True
    
    # Run pipeline
    pipeline = ResearchPipeline(data_dir=data_dir, models_dir=models_dir)
    
    if train_only:
        print("Training models with cached data...")
        combined_data = pipeline.process_datasets()
        await pipeline.train_models(combined_data)
        pipeline.generate_reports()
    elif fetch_only:
        print("Fetching datasets only...")
        await pipeline.fetch_datasets()
    else:
        await pipeline.run_full_pipeline()

if __name__ == "__main__":
    asyncio.run(main())