scripts/adapters/base_adapter.py

"""
Base adapter class for external dataset ingestion.

This module provides the abstract base class and common utilities for all
dataset adapters in ART-SafeBench v2.0.0.
"""

import json
import uuid
import hashlib
import logging
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Any, Optional, Generator
from dataclasses import dataclass, asdict

logger = logging.getLogger(__name__)

try:
    import jsonschema
    HAS_JSONSCHEMA = True
except ImportError:  # pragma: no cover
    HAS_JSONSCHEMA = False


@dataclass
class DatasetMetadata:
    """Metadata for external datasets."""
    name: str
    version: str
    license: str
    license_url: str
    source_url: str
    paper_url: Optional[str]
    can_redistribute: bool
    target_surface: str  # B1, B2, B3, or B4
    description: str
    record_count: Optional[int] = None
    last_updated: Optional[str] = None


@dataclass
class AdapterResult:
    """Result of an adapter transformation."""
    success: bool
    records_processed: int
    records_failed: int
    output_file: Optional[str]
    errors: List[str]
    warnings: List[str]


class BaseAdapter(ABC):
    """
    Abstract base class for dataset adapters.

    All dataset adapters must inherit from this class and implement
    the required methods for data transformation.
    """

    BENCHMARK_VERSION = "2.0.0"
    _SCHEMA_BY_SURFACE = {
        "B1": "schema/rag_poisoning.schema.json",
        "B2": "schema/image_poisoning.schema.json",
        "B3": "schema/safety_direct_query.schema.json",
        "B4": "schema/orchestrator_attacks.schema.json",
    }

    def __init__(self, output_dir: Path, verbose: bool = False, validate_schema: bool = True):
        """
        Initialize the adapter.

        Args:
            output_dir: Directory to write output files
            verbose: Enable verbose logging
            validate_schema: Validate transformed records against the benchmark JSON schema
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.verbose = verbose
        self.validate_schema = bool(validate_schema)
        self._setup_logging()
        self._schema_validator = None

        if self.validate_schema and not HAS_JSONSCHEMA:
            logger.warning("jsonschema not installed; disabling schema validation")
            self.validate_schema = False

    def _setup_logging(self) -> None:
        """Configure logging based on verbosity."""
        level = logging.DEBUG if self.verbose else logging.INFO
        logging.basicConfig(
            level=level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )

    @property
    @abstractmethod
    def metadata(self) -> DatasetMetadata:
        """Return metadata about the external dataset."""
        pass

    @abstractmethod
    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Transform a single record to the target schema.

        Args:
            record: Raw record from the external dataset

        Returns:
            Transformed record matching target schema, or None if invalid
        """
        pass

    @abstractmethod
    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        """
        Load and yield records from the source dataset.

        Yields:
            Raw records from the external dataset
        """
        pass

    def generate_id(self, prefix: str, content: str = None) -> str:
        """
        Generate a unique, reproducible ID for a record.

        Args:
            prefix: ID prefix (e.g., "B3-HarmBench")
            content: Optional content to hash for reproducibility

        Returns:
            Unique identifier string
        """
        if content:
            # Create deterministic ID based on content hash
            content_hash = hashlib.sha256(content.encode()).hexdigest()[:12]
            return f"{prefix}-{content_hash}"
        else:
            # Create random UUID-based ID
            return f"{prefix}-{uuid.uuid4().hex[:8]}"

    def get_timestamp(self) -> str:
        """Get current UTC timestamp in ISO format."""
        return datetime.now(timezone.utc).isoformat()

    def validate_record(self, record: Dict[str, Any], required_fields: List[str]) -> bool:
        """
        Validate that a record contains all required fields.

        Args:
            record: Record to validate
            required_fields: List of required field names

        Returns:
            True if valid, False otherwise
        """
        for field in required_fields:
            if field not in record or record[field] is None:
                logger.warning(f"Missing required field: {field}")
                return False
        return True

    def sanitize_text(self, text: str, max_length: int = None) -> str:
        """
        Sanitize text content for safe storage.

        Args:
            text: Text to sanitize
            max_length: Optional maximum length

        Returns:
            Sanitized text
        """
        if not isinstance(text, str):
            text = str(text) if text is not None else ""

        # Remove null bytes and other problematic characters
        text = text.replace('\x00', '')

        # Truncate if necessary
        if max_length and len(text) > max_length:
            text = text[:max_length] + "...[TRUNCATED]"

        return text

    def save_jsonl(self, records: List[Dict[str, Any]], filename: str) -> str:
        """
        Save records to a JSONL file.

        Args:
            records: List of records to save
            filename: Output filename

        Returns:
            Path to the saved file
        """
        output_path = self.output_dir / filename

        with open(output_path, 'w', encoding='utf-8') as f:
            for record in records:
                f.write(json.dumps(record, ensure_ascii=False) + '\n')

        logger.info(f"Saved {len(records)} records to {output_path}")
        return str(output_path)

    def _get_schema_path(self) -> Optional[Path]:
        """Resolve the JSON schema path for this adapter's target surface."""
        rel = self._SCHEMA_BY_SURFACE.get(self.metadata.target_surface)
        if not rel:
            return None
        # base_adapter.py -> scripts/adapters -> scripts -> repo root
        repo_root = Path(__file__).resolve().parents[2]
        return repo_root / rel

    def _get_schema_validator(self):
        """Lazily construct and cache the jsonschema validator."""
        if not self.validate_schema:
            return None
        if self._schema_validator is not None:
            return self._schema_validator

        schema_path = self._get_schema_path()
        if not schema_path or not schema_path.exists():
            logger.warning(
                "Schema file not found for %s (%s); disabling schema validation",
                self.metadata.name,
                self.metadata.target_surface,
            )
            self.validate_schema = False
            return None

        schema = json.loads(schema_path.read_text(encoding="utf-8"))
        self._schema_validator = jsonschema.Draft7Validator(schema)
        return self._schema_validator

    def run(self, max_records: int = None) -> AdapterResult:
        """
        Execute the full adapter pipeline.

        Args:
            max_records: Optional limit on number of records to process

        Returns:
            AdapterResult with processing statistics
        """
        logger.info(f"Starting adapter: {self.metadata.name}")
        logger.info(f"Target surface: {self.metadata.target_surface}")
        logger.info(f"License: {self.metadata.license}")

        records = []
        errors = []
        warnings = []
        processed = 0
        failed = 0

        try:
            validator = self._get_schema_validator()
            for raw_record in self.load_source_data():
                if max_records and processed >= max_records:
                    warnings.append(f"Stopped at max_records limit: {max_records}")
                    break

                try:
                    transformed = self.transform_record(raw_record)
                    if transformed:
                        if validator is not None:
                            schema_errors = list(validator.iter_errors(transformed))
                            if schema_errors:
                                failed += 1
                                first = schema_errors[0]
                                rec_id = (
                                    transformed.get("id")
                                    or transformed.get("record_id")
                                    or transformed.get("attack_id")
                                    or "<unknown>"
                                )
                                errors.append(
                                    f"Schema validation failed ({rec_id}): {first.message}"
                                )
                                continue
                        records.append(transformed)
                        processed += 1
                    else:
                        failed += 1
                except Exception as e:
                    failed += 1
                    errors.append(f"Transform error: {str(e)}")
                    if self.verbose:
                        logger.exception("Record transformation failed")

            # Save output
            if records:
                filename = f"{self.metadata.target_surface}_augmented_{self.metadata.name.lower().replace(' ', '_')}.jsonl"
                output_file = self.save_jsonl(records, filename)
            else:
                output_file = None
                warnings.append("No records were successfully transformed")

            return AdapterResult(
                success=len(records) > 0,
                records_processed=processed,
                records_failed=failed,
                output_file=output_file,
                errors=errors,
                warnings=warnings
            )

        except Exception as e:
            logger.exception("Adapter execution failed")
            return AdapterResult(
                success=False,
                records_processed=processed,
                records_failed=failed,
                output_file=None,
                errors=[f"Fatal error: {str(e)}"] + errors,
                warnings=warnings
            )