glitchlings 0.10.2__cp312-cp312-macosx_11_0_universal2.whl

glitchlings/attack/core_planning.py

@@ -0,0 +1,612 @@
+"""Pure planning functions for Attack orchestration.
+
+This module contains deterministic, side-effect-free logic for planning
+attack execution and assembling results. Functions here operate on
+already-resolved inputs without performing IO or invoking FFI.
+
+**Design Philosophy:**
+
+All functions in this module are *pure* - they perform planning and
+composition based solely on their inputs, without side effects. They do not:
+- Import or invoke Rust FFI
+- Resolve tokenizers or glitchlings
+- Create RNG instances
+- Perform I/O of any kind
+
+The separation allows:
+- Plan verification without Rust dependencies
+- Unit testing of orchestration logic in isolation
+- Clear boundaries between planning and execution
+
+See AGENTS.md "Functional Purity Architecture" for full details.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, TypeGuard
+
+from ..util.transcripts import Transcript
+
+# ---------------------------------------------------------------------------
+# Type Guards
+# ---------------------------------------------------------------------------
+
+
+def is_string_batch(value: Any) -> TypeGuard[Sequence[str]]:
+    """Determine if value is a batch of strings (not a single string).
+
+    Args:
+        value: Value to check.
+
+    Returns:
+        True if value is a non-string sequence of strings.
+    """
+    if isinstance(value, (str, bytes)):
+        return False
+    if not isinstance(value, Sequence):
+        return False
+    return all(isinstance(item, str) for item in value)
+
+
+def is_transcript_like(value: Any) -> bool:
+    """Check if value resembles a transcript structure.
+
+    A transcript is a sequence of mappings with 'role' and 'content' keys.
+
+    Args:
+        value: Value to check.
+
+    Returns:
+        True if value appears to be a transcript.
+    """
+    if not isinstance(value, Sequence):
+        return False
+    if isinstance(value, (str, bytes)):
+        return False
+    if not value:
+        return True  # Empty sequence could be empty transcript
+    first = value[0]
+    return isinstance(first, Mapping) and "content" in first
+
+
+# ---------------------------------------------------------------------------
+# Attack Planning
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class AttackPlan:
+    """Pure representation of what an Attack will do.
+
+    Attributes:
+        input_type: Type of input ("string", "batch", "transcript").
+        original_contents: List of content strings to process.
+        batch_size: Number of items in the batch.
+    """
+
+    input_type: str
+    original_contents: list[str]
+    batch_size: int
+
+    @property
+    def is_batch(self) -> bool:
+        """Return True if this plan represents batched input."""
+        return self.input_type in ("batch", "transcript")
+
+    @property
+    def is_empty(self) -> bool:
+        """Return True if there are no contents to process."""
+        return self.batch_size == 0
+
+
+def plan_attack(text: str | Transcript | Sequence[str]) -> AttackPlan:
+    """Create an execution plan for the given input.
+
+    This pure function analyzes the input structure and creates a plan
+    without actually executing anything.
+
+    Args:
+        text: Input text, transcript, or batch of strings.
+
+    Returns:
+        AttackPlan describing how to process the input.
+
+    Raises:
+        TypeError: If input type is not recognized.
+    """
+    if is_string_batch(text):
+        contents = list(text)
+        return AttackPlan(
+            input_type="batch",
+            original_contents=contents,
+            batch_size=len(contents),
+        )
+
+    if is_transcript_like(text):
+        contents = extract_transcript_contents(text)  # type: ignore[arg-type]
+        return AttackPlan(
+            input_type="transcript",
+            original_contents=contents,
+            batch_size=len(contents),
+        )
+
+    if isinstance(text, str):
+        return AttackPlan(
+            input_type="string",
+            original_contents=[text],
+            batch_size=1,
+        )
+
+    message = f"Attack expects string, transcript, or list of strings, got {type(text).__name__}"
+    raise TypeError(message)
+
+
+def extract_transcript_contents(transcript: Sequence[Mapping[str, Any]]) -> list[str]:
+    """Extract content strings from a transcript (pure version).
+
+    Args:
+        transcript: Sequence of turn mappings with 'content' keys.
+
+    Returns:
+        List of content strings.
+
+    Raises:
+        TypeError: If transcript structure is invalid.
+    """
+    contents: list[str] = []
+    for index, turn in enumerate(transcript):
+        if not isinstance(turn, Mapping):
+            raise TypeError(f"Transcript turn #{index + 1} must be a mapping.")
+        content = turn.get("content")
+        if not isinstance(content, str):
+            raise TypeError(f"Transcript turn #{index + 1} is missing string content.")
+        contents.append(content)
+    return contents
+
+
+# ---------------------------------------------------------------------------
+# Result Planning
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class MetricPlan:
+    """Plan for computing a single metric.
+
+    Attributes:
+        name: Metric name.
+        use_batch: Whether to use batch computation.
+    """
+
+    name: str
+    use_batch: bool
+
+
+@dataclass(frozen=True, slots=True)
+class ResultPlan:
+    """Plan for assembling attack results.
+
+    Attributes:
+        is_batch: Whether results are batched.
+        metric_names: Names of metrics to compute.
+        tokenizer_info: Description of tokenizer being used.
+    """
+
+    is_batch: bool
+    metric_names: tuple[str, ...]
+    tokenizer_info: str
+
+    def format_metrics(
+        self,
+        raw_metrics: dict[str, float | list[float]],
+    ) -> dict[str, float | list[float]]:
+        """Format metrics according to the result type.
+
+        For single results, collapses list metrics to scalars.
+        For batch results, ensures all metrics are lists.
+
+        Args:
+            raw_metrics: Raw metric values from computation.
+
+        Returns:
+            Formatted metrics appropriate for the result type.
+        """
+        if self.is_batch:
+            return _format_metrics_for_batch(raw_metrics)  # type: ignore[return-value]
+        return _format_metrics_for_single(raw_metrics)  # type: ignore[return-value]
+
+
+def plan_result(
+    attack_plan: AttackPlan,
+    metric_names: Sequence[str],
+    tokenizer_info: str,
+) -> ResultPlan:
+    """Create a plan for assembling results.
+
+    Args:
+        attack_plan: The attack execution plan.
+        metric_names: Names of metrics being computed.
+        tokenizer_info: Description of the tokenizer.
+
+    Returns:
+        ResultPlan for assembling the final result.
+    """
+    return ResultPlan(
+        is_batch=attack_plan.is_batch,
+        metric_names=tuple(metric_names),
+        tokenizer_info=tokenizer_info,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Metric Formatting (Pure)
+# ---------------------------------------------------------------------------
+
+
+def _format_metrics_for_single(
+    metrics: dict[str, float | list[float]],
+) -> dict[str, float]:
+    """Collapse batch metrics to single values.
+
+    Args:
+        metrics: Raw metrics that may be lists.
+
+    Returns:
+        Metrics as scalar floats.
+    """
+    result: dict[str, float] = {}
+    for name, value in metrics.items():
+        if isinstance(value, list):
+            result[name] = value[0] if value else 0.0
+        else:
+            result[name] = value
+    return result
+
+
+def _format_metrics_for_batch(
+    metrics: dict[str, float | list[float]],
+) -> dict[str, list[float]]:
+    """Ensure all metrics are lists.
+
+    Args:
+        metrics: Raw metrics that may be scalars.
+
+    Returns:
+        Metrics as lists of floats.
+    """
+    result: dict[str, list[float]] = {}
+    for name, value in metrics.items():
+        if isinstance(value, list):
+            result[name] = list(value)
+        else:
+            result[name] = [value]
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Batch Adapter (Pure)
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class BatchAdapter:
+    """Adapter that normalizes all inputs to batch format internally.
+
+    This adapter wraps single strings as batches of size 1, allowing
+    uniform processing throughout the attack pipeline. It tracks whether
+    to unwrap results back to single format at output time.
+
+    Attributes:
+        contents: List of content strings (always a list, even for single).
+        unwrap_single: True if the original input was a single string.
+        input_type: Original input type ("string", "batch", "transcript").
+    """
+
+    contents: list[str]
+    unwrap_single: bool
+    input_type: str
+
+    @classmethod
+    def from_plan(cls, plan: "AttackPlan") -> "BatchAdapter":
+        """Create a BatchAdapter from an AttackPlan.
+
+        Args:
+            plan: The attack execution plan.
+
+        Returns:
+            BatchAdapter configured for the plan's input type.
+        """
+        return cls(
+            contents=plan.original_contents,
+            unwrap_single=plan.input_type == "string",
+            input_type=plan.input_type,
+        )
+
+    def unwrap_tokens(
+        self,
+        tokens: list[list[str]],
+    ) -> list[str] | list[list[str]]:
+        """Unwrap batched tokens to match original input format.
+
+        Args:
+            tokens: Batched token lists (2D).
+
+        Returns:
+            1D list for single input, 2D list for batch input.
+        """
+        if self.unwrap_single and tokens:
+            return tokens[0]
+        return tokens
+
+    def unwrap_token_ids(
+        self,
+        token_ids: list[list[int]],
+    ) -> list[int] | list[list[int]]:
+        """Unwrap batched token IDs to match original input format.
+
+        Args:
+            token_ids: Batched token ID lists (2D).
+
+        Returns:
+            1D list for single input, 2D list for batch input.
+        """
+        if self.unwrap_single and token_ids:
+            return token_ids[0]
+        return token_ids
+
+    def unwrap_metrics(
+        self,
+        metrics: dict[str, list[float]],
+    ) -> dict[str, float | list[float]]:
+        """Unwrap batched metrics to match original input format.
+
+        Args:
+            metrics: Batched metrics (values are lists).
+
+        Returns:
+            Scalar metrics for single input, list metrics for batch.
+        """
+        if self.unwrap_single:
+            return {name: values[0] if values else 0.0 for name, values in metrics.items()}
+        # Explicitly construct new dict to satisfy type checker (dict invariance)
+        result: dict[str, float | list[float]] = {name: values for name, values in metrics.items()}
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Result Assembly (Pure)
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class EncodedData:
+    """Encoded token data for result assembly.
+
+    Tokens and IDs are always stored in batch format (2D lists)
+    internally. Use BatchAdapter.unwrap_* methods to convert to
+    the appropriate output format.
+
+    Attributes:
+        tokens: Token strings as batched 2D list.
+        token_ids: Token IDs as batched 2D list.
+    """
+
+    tokens: list[list[str]]
+    token_ids: list[list[int]]
+
+
+def assemble_single_result_fields(
+    *,
+    original: str,
+    corrupted: str,
+    input_encoded: EncodedData,
+    output_encoded: EncodedData,
+    tokenizer_info: str,
+    metrics: dict[str, float],
+) -> dict[str, object]:
+    """Assemble field dictionary for single-string AttackResult.
+
+    Args:
+        original: Original input string.
+        corrupted: Corrupted output string.
+        input_encoded: Encoded original tokens.
+        output_encoded: Encoded corrupted tokens.
+        tokenizer_info: Tokenizer description.
+        metrics: Computed metrics (scalar).
+
+    Returns:
+        Dictionary suitable for AttackResult construction.
+    """
+    # For single strings, tokens are batched internally as [[...]]
+    # so we unwrap the first (and only) element
+    input_tokens = input_encoded.tokens[0] if input_encoded.tokens else []
+    output_tokens = output_encoded.tokens[0] if output_encoded.tokens else []
+    input_ids = input_encoded.token_ids[0] if input_encoded.token_ids else []
+    output_ids = output_encoded.token_ids[0] if output_encoded.token_ids else []
+
+    return {
+        "original": original,
+        "corrupted": corrupted,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "input_token_ids": input_ids,
+        "output_token_ids": output_ids,
+        "tokenizer_info": tokenizer_info,
+        "metrics": metrics,
+    }
+
+
+def assemble_batch_result_fields(
+    *,
+    original: Transcript | Sequence[str],
+    corrupted: Transcript | Sequence[str],
+    input_encoded: EncodedData,
+    output_encoded: EncodedData,
+    tokenizer_info: str,
+    metrics: dict[str, float | list[float]],
+) -> dict[str, object]:
+    """Assemble field dictionary for batched AttackResult.
+
+    Args:
+        original: Original transcript or string batch.
+        corrupted: Corrupted transcript or string batch.
+        input_encoded: Encoded original tokens (batched).
+        output_encoded: Encoded corrupted tokens (batched).
+        tokenizer_info: Tokenizer description.
+        metrics: Computed metrics (list per batch item).
+
+    Returns:
+        Dictionary suitable for AttackResult construction.
+    """
+    return {
+        "original": original,
+        "corrupted": corrupted,
+        "input_tokens": input_encoded.tokens,
+        "output_tokens": output_encoded.tokens,
+        "input_token_ids": input_encoded.token_ids,
+        "output_token_ids": output_encoded.token_ids,
+        "tokenizer_info": tokenizer_info,
+        "metrics": metrics,
+    }
+
+
+def assemble_empty_result_fields(
+    *,
+    original: Transcript | Sequence[str],
+    corrupted: Transcript | Sequence[str],
+    tokenizer_info: str,
+    metric_names: Sequence[str],
+) -> dict[str, object]:
+    """Assemble field dictionary for empty input.
+
+    Args:
+        original: Original empty transcript or list.
+        corrupted: Corrupted empty transcript or list.
+        tokenizer_info: Tokenizer description.
+        metric_names: Names of metrics to include as empty.
+
+    Returns:
+        Dictionary suitable for AttackResult construction.
+    """
+    return {
+        "original": original,
+        "corrupted": corrupted,
+        "input_tokens": [],
+        "output_tokens": [],
+        "input_token_ids": [],
+        "output_token_ids": [],
+        "tokenizer_info": tokenizer_info,
+        "metrics": {name: [] for name in metric_names},
+    }
+
+
+def assemble_result_fields(
+    *,
+    adapter: BatchAdapter,
+    original: str | Transcript | Sequence[str],
+    corrupted: str | Transcript | Sequence[str],
+    input_encoded: EncodedData,
+    output_encoded: EncodedData,
+    tokenizer_info: str,
+    metrics: dict[str, list[float]],
+) -> dict[str, object]:
+    """Assemble AttackResult fields using batch adapter for uniform handling.
+
+    This function uses the BatchAdapter to handle both single and batch
+    inputs uniformly. Internally, all data is processed as batches, then
+    unwrapped appropriately based on the original input type.
+
+    Args:
+        adapter: BatchAdapter tracking input type.
+        original: Original input container.
+        corrupted: Corrupted output container.
+        input_encoded: Encoded original tokens (always batched internally).
+        output_encoded: Encoded corrupted tokens (always batched internally).
+        tokenizer_info: Tokenizer description.
+        metrics: Computed metrics (always batched as lists internally).
+
+    Returns:
+        Dictionary suitable for AttackResult construction.
+    """
+    return {
+        "original": original,
+        "corrupted": corrupted,
+        "input_tokens": adapter.unwrap_tokens(input_encoded.tokens),
+        "output_tokens": adapter.unwrap_tokens(output_encoded.tokens),
+        "input_token_ids": adapter.unwrap_token_ids(input_encoded.token_ids),
+        "output_token_ids": adapter.unwrap_token_ids(output_encoded.token_ids),
+        "tokenizer_info": tokenizer_info,
+        "metrics": adapter.unwrap_metrics(metrics),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Token Count Helpers (Pure)
+# ---------------------------------------------------------------------------
+
+
+def compute_token_counts(
+    input_tokens: list[str] | list[list[str]],
+    output_tokens: list[str] | list[list[str]],
+) -> tuple[list[int], list[int]]:
+    """Compute token counts for inputs and outputs.
+
+    Handles both single sequences and batches.
+
+    Args:
+        input_tokens: Input token sequence(s).
+        output_tokens: Output token sequence(s).
+
+    Returns:
+        Tuple of (input_counts, output_counts) as lists.
+    """
+    # Check if batched
+    if input_tokens and isinstance(input_tokens[0], list):
+        input_counts = [len(batch) for batch in input_tokens]
+        output_counts = [len(batch) for batch in output_tokens]
+    else:
+        input_counts = [len(input_tokens)]
+        output_counts = [len(output_tokens)]
+    return input_counts, output_counts
+
+
+def format_token_count_delta(input_count: int, output_count: int) -> str:
+    """Format a token count change as a string.
+
+    Args:
+        input_count: Number of input tokens.
+        output_count: Number of output tokens.
+
+    Returns:
+        Formatted string like "10 -> 12 (+2)".
+    """
+    delta = output_count - input_count
+    return f"{input_count} -> {output_count} ({delta:+d})"
+
+
+__all__ = [
+    # Type guards
+    "is_string_batch",
+    "is_transcript_like",
+    # Attack planning
+    "AttackPlan",
+    "plan_attack",
+    "extract_transcript_contents",
+    # Result planning
+    "MetricPlan",
+    "ResultPlan",
+    "plan_result",
+    # Batch adapter
+    "BatchAdapter",
+    # Result assembly
+    "EncodedData",
+    "assemble_result_fields",
+    "assemble_single_result_fields",
+    "assemble_batch_result_fields",
+    "assemble_empty_result_fields",
+    # Token counts
+    "compute_token_counts",
+    "format_token_count_delta",
+]