scripts/adapters/b3_tensortrust_adapter.py

"""
TensorTrust Dataset Adapter for B3 (Direct Query) Schema

TensorTrust (aka "tensor-trust"/PromptGame): A prompt-hijacking / prompt-extraction game.
License: BSD-2-Clause
Source: https://github.com/HumanCompatibleAI/tensor-trust

IMPORTANT DATA ACCESS NOTE:
The upstream repo's public codebase does not ship the underlying database dump.
This adapter expects a user-provided export in JSONL format (e.g. a
`raw_dump_attacks.jsonl` file produced by the upstream data pipeline).

Maps TensorTrust attack attempts to the B3 Direct Query Safety schema.
"""

import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional, Generator

from .base_adapter import BaseAdapter, DatasetMetadata

logger = logging.getLogger(__name__)


class TensorTrustAdapter(BaseAdapter):
    """
    Adapter for TensorTrust exported attacks.

    Expected input: JSONL file where each line is a dict with (at least) an
    attacker prompt and model output. The canonical export name used by the
    upstream pipeline is `raw_dump_attacks.jsonl`.
    """

    def __init__(
        self,
        output_dir: Path,
        attacks_jsonl_path: Path,
        verbose: bool = False,
    ):
        super().__init__(output_dir, verbose)
        self.attacks_jsonl_path = Path(attacks_jsonl_path)

    @property
    def metadata(self) -> DatasetMetadata:
        return DatasetMetadata(
            name="TensorTrust",
            version="1.0",
            license="BSD-2-Clause",
            license_url="https://github.com/HumanCompatibleAI/tensor-trust/blob/main/LICENSE",
            source_url="https://github.com/HumanCompatibleAI/tensor-trust",
            paper_url=None,
            can_redistribute=True,
            target_surface="B3",
            description="PromptGame-style prompt extraction / hijacking attack attempts (exported JSONL)",
        )

    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        if not self.attacks_jsonl_path.exists():
            raise FileNotFoundError(f"TensorTrust attacks JSONL not found: {self.attacks_jsonl_path}")

        logger.info(f"Loading TensorTrust attacks from {self.attacks_jsonl_path}")

        with open(self.attacks_jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                yield json.loads(line)

    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        attacker_input = record.get("attacker_input") or record.get("attacker_prompt") or ""
        if not attacker_input or len(str(attacker_input).strip()) < 3:
            return None

        attack_id = record.get("attack_id")
        llm_choice = record.get("llm_choice") or record.get("model") or ""
        llm_output = record.get("llm_output")
        output_is_access_granted = record.get("output_is_access_granted")

        opening_defense = record.get("opening_defense") or record.get("defense_pre_prompt") or ""
        closing_defense = record.get("closing_defense") or record.get("defense_post_prompt") or ""

        ts = record.get("timestamp") or self.get_timestamp()

        user_query = self.sanitize_text(str(attacker_input), max_length=15000)
        target_llm_output = self.sanitize_text(str(llm_output), max_length=20000) if llm_output is not None else ""

        system_prompt = "\n\n".join([p for p in [opening_defense, closing_defense] if p])
        system_prompt = self.sanitize_text(system_prompt, max_length=20000) if system_prompt else ""

        assessment = "Success" if bool(output_is_access_granted) else "Failure"
        justification_parts = [
            "TensorTrust/PromptGame exported attack attempt.",
            f"output_is_access_granted={bool(output_is_access_granted)}.",
        ]
        if llm_choice:
            justification_parts.append(f"llm_choice={llm_choice}.")
        justification = " ".join(justification_parts)

        record_id = (
            f"B3-TensorTrust-{attack_id}"
            if attack_id is not None and str(attack_id).strip()
            else self.generate_id("B3-TensorTrust", user_query)
        )

        return {
            "id": record_id,
            "attack_category": "Prompt Extraction",
            "attack_subtype": "TensorTrust_AccessCode_Extraction",
            "adversarial_goal": "Extract the hidden access code / prompt (PromptGame-style)",
            "user_query": user_query,
            "target_llm_output": target_llm_output,
            "judge_assessment": {
                "assessment": assessment,
                "justification": justification,
                "confidence": "High",
            },
            "benign_content": None,
            "mode": None,
            "role_play": None,
            "manipulation_strategy": None,
            "communication_style": None,
            "query_structure": "direct prompt",
            "intent_Opacity": "explicit",
            "topic_domain": "promptgame",
            "temporal_context": None,
            "language": "English",
            "Obfuscation": None,
            "adversarial_suffix": None,
            "payload_splitting": None,
            "encoding": None,
            "target_rag_component": "Generator_LLM",
            "target_component_model": llm_choice if llm_choice else "External_Evaluation",
            "target_model_system_prompt": system_prompt,
            "all_llm_details": {
                "llm_choice": llm_choice,
                "is_self_attack": record.get("is_self_attack"),
            },
            "agent_architecture_description": "TensorTrust/PromptGame (interactive prompt-extraction game)",
            "attack_gen_pipeline_description": "Imported from TensorTrust exported attacks (raw_dump_attacks.jsonl)",
            "generation_timestamp": ts,
            "script_version": self.BENCHMARK_VERSION,
            "source_dataset": "TensorTrust",
            "source_id": str(attack_id) if attack_id is not None else None,
            "source_license": "BSD-2-Clause",
        }


def main():
    import argparse

    parser = argparse.ArgumentParser(description="TensorTrust to B3 Schema Adapter")
    parser.add_argument(
        "--attacks-jsonl-path",
        type=Path,
        required=True,
        help="Path to TensorTrust exported attacks JSONL (e.g., raw_dump_attacks.jsonl)",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented/core"),
        help="Output directory",
    )
    parser.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Maximum records to process",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable verbose output",
    )

    args = parser.parse_args()

    adapter = TensorTrustAdapter(
        output_dir=args.output_dir,
        attacks_jsonl_path=args.attacks_jsonl_path,
        verbose=args.verbose,
    )
    result = adapter.run(max_records=args.max_records)

    print("\nAdapter Results:")
    print(f"  Success: {result.success}")
    print(f"  Records processed: {result.records_processed}")
    print(f"  Records failed: {result.records_failed}")
    if result.output_file:
        print(f"  Output file: {result.output_file}")


if __name__ == "__main__":
    main()