diff --git a/hackagent/attacks/base.py b/hackagent/attacks/base.py index 4c01b461..2615572e 100644 --- a/hackagent/attacks/base.py +++ b/hackagent/attacks/base.py @@ -2,7 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import abc -from typing import Any, Dict +from typing import Any, Dict, List + +from hackagent.attacks.types import AttackResult class BaseAttack(abc.ABC): @@ -64,7 +66,7 @@ def _setup(self): pass @abc.abstractmethod - def run(self, **kwargs: Any) -> Any: + def run(self, **kwargs: Any) -> List[AttackResult]: """ Executes the attack logic. @@ -80,12 +82,7 @@ def run(self, **kwargs: Any) -> Any: - target_model: The model to attack Returns: - Attack-specific results. The format varies by implementation but - typically includes: - - adversarial_examples: Generated adversarial inputs - - success_metrics: Attack success rates and statistics - - detailed_results: Comprehensive result data (e.g., pandas DataFrame) - - attack_report: Summary of attack performance + A list of :class:`~hackagent.attacks.types.AttackResult` instances. Raises: NotImplementedError: If the method is not implemented by a subclass. diff --git a/hackagent/attacks/orchestrator.py b/hackagent/attacks/orchestrator.py index 088855b2..8ff1fe8a 100644 --- a/hackagent/attacks/orchestrator.py +++ b/hackagent/attacks/orchestrator.py @@ -56,6 +56,11 @@ DEFAULT_REMOTE_ROLE_ENDPOINT, ) from hackagent.server.storage.enums import StatusEnum +from hackagent.attacks.types import ( + AttackResult, + attack_results_to_rows, + flatten_run_result, +) if TYPE_CHECKING: from hackagent.agent import HackAgent @@ -1446,24 +1451,6 @@ def _get_attack_impl_kwargs( "agent_router": agent_router, } - @staticmethod - def _normalize_attack_results(results: Any) -> List[Dict[str, Any]]: - """Normalize heterogeneous attack outputs into a list of row dicts.""" - if results is None: - return [] - if isinstance(results, list): - return results - if isinstance(results, dict): - evaluated = results.get("evaluated") - if isinstance(evaluated, list): - return evaluated - for key in ("rows", "results", "data", "items"): - value = results.get(key) - if isinstance(value, list): - return value - return [] - return [] - def _execute_local_attack( self, attack_id: str, @@ -1471,7 +1458,7 @@ def _execute_local_attack( attack_params: Dict[str, Any], attack_config: Dict[str, Any], run_config_override: Optional[Dict[str, Any]], - ) -> Any: + ) -> List[AttackResult]: """ Execute attack locally using technique implementation. @@ -1560,7 +1547,7 @@ def _execute_local_attack( f"goal_batch_workers={goal_batch_workers} (parallel goals per batch)" ) - all_results: List[Dict[str, Any]] = [] + all_results: List[AttackResult] = [] batch_timings: List[float] = [] for batch_idx, (batch_start_idx, batch_goals) in enumerate(batches): @@ -1576,11 +1563,7 @@ def _execute_local_attack( # Global run status is finalized once in execute(). attack_impl.config["_suppress_run_status_updates"] = True batch_params = {**attack_params, "goals": batch_goals} - # attack_impl.run() may return a dict (e.g. baseline's - # {"evaluated": [...], "summary": [...]}) rather than a - # flat row list — normalize before aggregating, otherwise - # extend() below would iterate the dict's *keys*. - batch_results = self._normalize_attack_results( + batch_results = flatten_run_result( attack_impl.run(**batch_params) ) else: @@ -1591,7 +1574,7 @@ def _run_single_goal( goal_idx_goal: Tuple[int, str], _batch_label: str = batch_label, _batch_start_idx: int = batch_start_idx, - ) -> Tuple[int, List[Dict[str, Any]]]: + ) -> Tuple[int, List[AttackResult]]: goal_idx, goal = goal_idx_goal # Label thread for _BatchContextFilter @@ -1614,16 +1597,14 @@ def _run_single_goal( } local_impl = self.attack_impl_class(**local_impl_kwargs) goal_params = {**attack_params, "goals": [goal]} - # Normalize here too — same dict-vs-list return shape - # concern as the sequential path above. - goal_results = self._normalize_attack_results( + goal_results = flatten_run_result( local_impl.run(**goal_params) ) logger.info(f"Goal done ({len(goal_results)} results)") return goal_idx, goal_results - per_goal_results: Dict[int, List[Dict[str, Any]]] = {} + per_goal_results: Dict[int, List[AttackResult]] = {} # Install a LogRecordFactory so *all* log records, # regardless of logger/handler routing, get the batch @@ -1677,7 +1658,7 @@ def _batch_record_factory(*args, **kwargs): ) return all_results - results = attack_impl.run(**attack_params) + results = flatten_run_result(attack_impl.run(**attack_params)) logger.info(f"{self.attack_type} attack completed") return results finally: @@ -1870,7 +1851,7 @@ def execute( attack_config=attack_config, run_config_override=effective_run_config, ) - normalized_results = self._normalize_attack_results(results) + normalized_results = attack_results_to_rows(results) # ========================= # RUN EVALUATION PIPELINE @@ -1904,10 +1885,11 @@ def execute( # (setting is_success/best_score per row) rather than via # a separate judge pass, so it only needs the shared # post-processing (default-filling + sync/ASR logging). - final_results = evaluator._postprocess_inline_judge_results( - normalized_results, attack_label="PAIR" + final_results = flatten_run_result( + evaluator._postprocess_inline_judge_results( + normalized_results, attack_label="PAIR" + ) ) - final_results = self._normalize_attack_results(final_results) evaluator.prepare_and_sync(final_results, run_id) logger.info("PAIR judge evaluation pipeline completed") else: @@ -1924,8 +1906,9 @@ def execute( ) # Run evaluation pipeline - final_results = evaluator.run_full_evaluation(normalized_results) - final_results = self._normalize_attack_results(final_results) + final_results = flatten_run_result( + evaluator.run_full_evaluation(normalized_results) + ) # Sync metrics to backend evaluator.prepare_and_sync(final_results, run_id) @@ -1942,7 +1925,7 @@ def execute( error=e, logger=logger, ) - final_results = results # fallback + final_results = normalized_results # fallback if _tui_event_bus is not None: _tui_event_bus.emit( "step_ended", diff --git a/hackagent/attacks/techniques/advprefix/attack.py b/hackagent/attacks/techniques/advprefix/attack.py index c83ed51b..b6de4f1c 100644 --- a/hackagent/attacks/techniques/advprefix/attack.py +++ b/hackagent/attacks/techniques/advprefix/attack.py @@ -20,6 +20,7 @@ from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging # Import step execution functions from same package @@ -269,7 +270,7 @@ def _get_pipeline_steps(self): ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Executes the full prefix generation pipeline. @@ -285,6 +286,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of dictionaries containing the final selected prefixes, or empty list if no prefixes were generated. """ + goals = goals or [] if not goals: return [] @@ -343,7 +345,7 @@ def run(self, goals: List[str]) -> List[Dict]: # Finalize pipeline-level tracking coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: # Crash-safe: mark all unfinalized goals as failed diff --git a/hackagent/attacks/techniques/autodan_turbo/attack.py b/hackagent/attacks/techniques/autodan_turbo/attack.py index e63cd784..74fb9cd2 100644 --- a/hackagent/attacks/techniques/autodan_turbo/attack.py +++ b/hackagent/attacks/techniques/autodan_turbo/attack.py @@ -5,10 +5,11 @@ import copy import logging import os -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from . import autodan_eval as evaluation, lifelong, warm_up from .config import DEFAULT_AUTODAN_TURBO_CONFIG, AutoDANTurboConfig @@ -152,7 +153,7 @@ def _get_pipeline_steps(self): return [] # Managed manually in run() (like PAIR) @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict[str, Any]]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """Execute full AutoDAN-Turbo pipeline. Pipeline mapping to paper/integration: @@ -169,6 +170,7 @@ def run(self, goals: List[str]) -> List[Dict[str, Any]]: Raises: Exception: Re-raises any runtime failure after coordinator finalization. """ + goals = goals or [] if not goals: return [] @@ -379,7 +381,7 @@ def run(self, goals: List[str]) -> List[Dict[str, Any]]: os.makedirs(output_dir, exist_ok=True) strategy_lib.save(f"{output_dir}/strategy_library") - return results + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("AutoDAN-Turbo failed") diff --git a/hackagent/attacks/techniques/base.py b/hackagent/attacks/techniques/base.py index 9ab8ec87..249316a1 100644 --- a/hackagent/attacks/techniques/base.py +++ b/hackagent/attacks/techniques/base.py @@ -31,6 +31,7 @@ from hackagent.logger import get_logger from typing import Any, Dict, List, Optional +from hackagent.attacks.types import AttackResult from hackagent.router.tracking import StepTracker, TrackingCoordinator logger = get_logger(__name__) @@ -451,7 +452,7 @@ def _get_pipeline_steps(self) -> List[Dict]: pass @abc.abstractmethod - def run(self, **kwargs) -> Any: + def run(self, **kwargs) -> List[AttackResult]: """ Execute the attack technique. @@ -460,12 +461,12 @@ def run(self, **kwargs) -> Any: 2. Define pipeline with self._get_pipeline_steps() 3. Execute pipeline with self._execute_pipeline() 4. Finalize with coordinator.finalize_all_goals() and coordinator.finalize_pipeline() - 5. Return results + 5. Return results as ``list[AttackResult]`` Args: **kwargs: Technique-specific parameters (e.g., goals, prompts) Returns: - Attack results (format varies by implementation) + A list of :class:`~hackagent.attacks.types.AttackResult` instances. """ pass diff --git a/hackagent/attacks/techniques/baseline/attack.py b/hackagent/attacks/techniques/baseline/attack.py index 25f7399d..6392239d 100644 --- a/hackagent/attacks/techniques/baseline/attack.py +++ b/hackagent/attacks/techniques/baseline/attack.py @@ -16,6 +16,7 @@ from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack from hackagent.attacks.shared.tui import with_tui_logging +from hackagent.attacks.types import AttackResult, rows_to_attack_results from . import generation from .config import DEFAULT_BASELINE_CONFIG @@ -159,7 +160,7 @@ def _build_step_args( return args @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> Dict[str, Any]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute baseline attack (direct goal submission). @@ -167,10 +168,11 @@ def run(self, goals: List[str]) -> Dict[str, Any]: goals: List of goal strings to send directly. Returns: - Dictionary with 'evaluated' and 'summary' DataFrames. + A list of :class:`~hackagent.attacks.types.AttackResult` instances. """ + goals = goals or [] if not goals: - return {"evaluated": [], "summary": []} + return [] coordinator = self._initialize_coordinator( attack_type="Baseline", @@ -187,7 +189,9 @@ def success_check(output): return output and isinstance(output, dict) coordinator.finalize_pipeline(results, success_check) - return results if results else {"evaluated": [], "summary": []} + return rows_to_attack_results( + results if results else {"evaluated": [], "summary": []} + ) except Exception as e: self.logger.error(f"Pipeline failed: {e}", exc_info=True) diff --git a/hackagent/attacks/techniques/bon/attack.py b/hackagent/attacks/techniques/bon/attack.py index 7050c43d..9612f7b1 100644 --- a/hackagent/attacks/techniques/bon/attack.py +++ b/hackagent/attacks/techniques/bon/attack.py @@ -39,6 +39,7 @@ from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.evaluator.evaluation_step import BaseEvaluationStep @@ -206,7 +207,7 @@ def _get_pipeline_steps(self) -> List[Dict]: # ------------------------------------------------------------------ @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """Execute the full BoN attack pipeline. The generation step performs the multi-step BoN search **and** inline @@ -220,6 +221,7 @@ def run(self, goals: List[str]) -> List[Dict]: Returns: List of result dictionaries, or empty list if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -273,7 +275,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("BoN pipeline failed with exception") diff --git a/hackagent/attacks/techniques/cipherchat/attack.py b/hackagent/attacks/techniques/cipherchat/attack.py index 929c3ac1..f6bd70ef 100644 --- a/hackagent/attacks/techniques/cipherchat/attack.py +++ b/hackagent/attacks/techniques/cipherchat/attack.py @@ -16,6 +16,7 @@ from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter @@ -160,7 +161,8 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: + goals = goals or [] if not goals: return [] @@ -203,7 +205,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.finalize_all_goals(results) coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("CipherChat pipeline failed with exception") diff --git a/hackagent/attacks/techniques/fc/attack.py b/hackagent/attacks/techniques/fc/attack.py index 354655a3..324e401e 100644 --- a/hackagent/attacks/techniques/fc/attack.py +++ b/hackagent/attacks/techniques/fc/attack.py @@ -24,6 +24,7 @@ from typing import Any, Dict, List, Optional from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.config import DEFAULT_JUDGE_IDENTIFIER from hackagent.router.router import AgentRouter @@ -250,7 +251,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the full FC-Attack pipeline. @@ -261,6 +262,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of dictionaries containing evaluation results, or empty list if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -292,7 +294,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("FC-Attack pipeline failed with exception") @@ -426,7 +428,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the full text-only flowchart attack pipeline. @@ -437,6 +439,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of dictionaries containing evaluation results, or empty list if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -469,7 +472,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("tFC pipeline failed with exception") diff --git a/hackagent/attacks/techniques/flipattack/attack.py b/hackagent/attacks/techniques/flipattack/attack.py index 33922dcd..041d7fe2 100644 --- a/hackagent/attacks/techniques/flipattack/attack.py +++ b/hackagent/attacks/techniques/flipattack/attack.py @@ -40,6 +40,7 @@ from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.config import DEFAULT_JUDGE_IDENTIFIER from hackagent.attacks.evaluator.evaluation_step import BaseEvaluationStep @@ -442,7 +443,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the full FlipAttack pipeline. @@ -457,6 +458,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of dictionaries containing evaluation results, or empty list if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -503,7 +505,7 @@ def run(self, goals: List[str]) -> List[Dict]: # Finalize pipeline-level tracking coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: # Crash-safe: mark all unfinalized goals as failed diff --git a/hackagent/attacks/techniques/h4rm3l/attack.py b/hackagent/attacks/techniques/h4rm3l/attack.py index f8c5cbad..628f7785 100644 --- a/hackagent/attacks/techniques/h4rm3l/attack.py +++ b/hackagent/attacks/techniques/h4rm3l/attack.py @@ -25,6 +25,7 @@ from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.evaluator.evaluation_step import BaseEvaluationStep @@ -242,7 +243,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the full h4rm3l attack pipeline. @@ -253,6 +254,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of result dicts with evaluation scores, or ``[]`` if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -301,7 +303,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("h4rm3l pipeline failed with exception") diff --git a/hackagent/attacks/techniques/indirect_prompt_injection/attack.py b/hackagent/attacks/techniques/indirect_prompt_injection/attack.py index 6b947587..ecc640c7 100644 --- a/hackagent/attacks/techniques/indirect_prompt_injection/attack.py +++ b/hackagent/attacks/techniques/indirect_prompt_injection/attack.py @@ -27,6 +27,7 @@ import numpy as np from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.router_factory import create_router from hackagent.attacks.shared.response_utils import extract_response_content from hackagent.config import DEFAULT_EMBEDDER_OPENAI_ENDPOINT @@ -360,7 +361,7 @@ def _get_rag_injection_params(self) -> Dict[str, Any]: params = self.config.get("rag_injection_params", {}) return params if isinstance(params, dict) else {} - def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[Dict[str, Any]]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the indirect prompt injection attack. @@ -460,7 +461,7 @@ def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[Dict[str, Any if not self.config.get("_suppress_run_status_updates", False): coordinator.finalize_pipeline(all_results) - return all_results + return rows_to_attack_results(all_results) def _run_single_goal( self, diff --git a/hackagent/attacks/techniques/mml/attack.py b/hackagent/attacks/techniques/mml/attack.py index 29a663f0..a8f69a7a 100644 --- a/hackagent/attacks/techniques/mml/attack.py +++ b/hackagent/attacks/techniques/mml/attack.py @@ -26,6 +26,7 @@ from typing import Any, Dict, List, Optional from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.config import DEFAULT_JUDGE_IDENTIFIER from hackagent.router.router import AgentRouter @@ -281,7 +282,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the full MML attack pipeline. @@ -295,6 +296,7 @@ def run(self, goals: List[str]) -> List[Dict]: List of dictionaries containing evaluation results, or empty list if no goals provided. """ + goals = goals or [] if not goals: return [] @@ -335,7 +337,7 @@ def run(self, goals: List[str]) -> List[Dict]: # Finalize pipeline-level tracking coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: # Crash-safe: mark all unfinalized goals as failed diff --git a/hackagent/attacks/techniques/pair/attack.py b/hackagent/attacks/techniques/pair/attack.py index 4481bb39..cdea3eba 100644 --- a/hackagent/attacks/techniques/pair/attack.py +++ b/hackagent/attacks/techniques/pair/attack.py @@ -24,6 +24,7 @@ from typing import Any, Dict, List, Optional from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.techniques.autodan_turbo.core import score_response from hackagent.attacks.techniques.config import ( DEFAULT_ATTACKER_IDENTIFIER, @@ -1189,7 +1190,7 @@ def _run_stream(stream_item: tuple[int, Dict[str, Any]]) -> bool: } @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict[str, Any]]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute PAIR attack on goals. @@ -1202,6 +1203,7 @@ def run(self, goals: List[str]) -> List[Dict[str, Any]]: Returns: List of attack results with scores """ + goals = goals or [] if not goals: return [] @@ -1330,7 +1332,7 @@ def _run_goal(i_goal: tuple) -> None: # Log summary via coordinator coordinator.log_summary() - return results + return rows_to_attack_results(results) except Exception as e: self.logger.error(f"PAIR attack failed: {e}", exc_info=True) diff --git a/hackagent/attacks/techniques/pap/attack.py b/hackagent/attacks/techniques/pap/attack.py index 06f59b1e..1dbc9530 100644 --- a/hackagent/attacks/techniques/pap/attack.py +++ b/hackagent/attacks/techniques/pap/attack.py @@ -26,6 +26,7 @@ from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.evaluator.evaluation_step import BaseEvaluationStep @@ -170,7 +171,7 @@ def _get_pipeline_steps(self) -> List[Dict]: # ------------------------------------------------------------------ @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """Execute the full PAP attack pipeline. Args: @@ -179,6 +180,7 @@ def run(self, goals: List[str]) -> List[Dict]: Returns: List of result dictionaries. """ + goals = goals or [] if not goals: return [] @@ -223,7 +225,7 @@ def run(self, goals: List[str]) -> List[Dict]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception: coordinator.finalize_on_error("PAP pipeline failed with exception") diff --git a/hackagent/attacks/techniques/rag/attack.py b/hackagent/attacks/techniques/rag/attack.py index 339b3147..5bbeb2d0 100644 --- a/hackagent/attacks/techniques/rag/attack.py +++ b/hackagent/attacks/techniques/rag/attack.py @@ -27,6 +27,7 @@ import numpy as np from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.attacks.shared.router_factory import create_router from hackagent.attacks.shared.response_utils import extract_response_content from hackagent.router.router import AgentRouter @@ -357,7 +358,7 @@ def _get_rag_injection_params(self) -> Dict[str, Any]: params = self.config.get("rag_injection_params", {}) return params if isinstance(params, dict) else {} - def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[Dict[str, Any]]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute the RAG Attack (indirect prompt injection). @@ -455,7 +456,7 @@ def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[Dict[str, Any if not self.config.get("_suppress_run_status_updates", False): coordinator.finalize_pipeline(all_results) - return all_results + return rows_to_attack_results(all_results) def _run_single_goal( self, diff --git a/hackagent/attacks/techniques/static_template/attack.py b/hackagent/attacks/techniques/static_template/attack.py index 4ca5b04b..6c8c1a84 100644 --- a/hackagent/attacks/techniques/static_template/attack.py +++ b/hackagent/attacks/techniques/static_template/attack.py @@ -16,6 +16,7 @@ from hackagent.router.router import AgentRouter from hackagent.attacks.techniques.base import BaseAttack from hackagent.attacks.shared.tui import with_tui_logging +from hackagent.attacks.types import AttackResult, rows_to_attack_results from . import generation, static_eval as evaluation from .config import DEFAULT_TEMPLATE_CONFIG @@ -228,7 +229,7 @@ def _build_step_args( return args @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> Dict[str, Any]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Execute static template attack. @@ -238,10 +239,11 @@ def run(self, goals: List[str]) -> Dict[str, Any]: goals: List of harmful goals to test Returns: - Dictionary with 'evaluated' and 'summary' DataFrames + A list of :class:`~hackagent.attacks.types.AttackResult` instances. """ + goals = goals or [] if not goals: - return {"evaluated": [], "summary": []} + return [] # Initialize unified coordinator coordinator = self._initialize_coordinator( @@ -264,7 +266,9 @@ def success_check(output): # Finalize pipeline-level tracking via coordinator coordinator.finalize_pipeline(results, success_check) - return results if results else {"evaluated": [], "summary": []} + return rows_to_attack_results( + results if results else {"evaluated": [], "summary": []} + ) except Exception as e: self.logger.error(f"Pipeline failed: {e}", exc_info=True) diff --git a/hackagent/attacks/techniques/tap/attack.py b/hackagent/attacks/techniques/tap/attack.py index bddebc74..a505284a 100644 --- a/hackagent/attacks/techniques/tap/attack.py +++ b/hackagent/attacks/techniques/tap/attack.py @@ -50,6 +50,7 @@ from hackagent.attacks.shared.tui import with_tui_logging from hackagent.attacks.techniques.base import BaseAttack +from hackagent.attacks.types import AttackResult, rows_to_attack_results from hackagent.server.client import AuthenticatedClient from hackagent.router.router import AgentRouter @@ -279,7 +280,7 @@ def _get_pipeline_steps(self) -> List[Dict]: ] @with_tui_logging(logger_name="hackagent.attacks", level=logging.INFO) - def run(self, goals: List[str]) -> List[Dict[str, Any]]: + def run(self, goals: Optional[List[str]] = None, **kwargs) -> List[AttackResult]: """ Run TAP end-to-end with unified tracking and pipeline steps. @@ -289,6 +290,7 @@ def run(self, goals: List[str]) -> List[Dict[str, Any]]: Returns: List of per-goal result dicts produced by the pipeline. """ + goals = goals or [] if not goals: return [] @@ -327,7 +329,7 @@ def run(self, goals: List[str]) -> List[Dict[str, Any]]: coordinator.log_summary() coordinator.finalize_pipeline(results) - return results if results is not None else [] + return rows_to_attack_results(results) except Exception as exc: self.logger.error(f"TAP attack failed: {exc}", exc_info=True) diff --git a/hackagent/attacks/types.py b/hackagent/attacks/types.py new file mode 100644 index 00000000..2ec30023 --- /dev/null +++ b/hackagent/attacks/types.py @@ -0,0 +1,190 @@ +# Copyright 2026 - AI4I. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed models for attack technique results. + +This module replaces the historical ``_normalize_attack_results()`` +duck-typing helper in :mod:`hackagent.attacks.orchestrator`, which used to +flatten heterogeneous technique outputs by probing for ``.evaluated``, +``.rows``, ``.results``, ``.data``, ``.items`` in turn. Any new technique +naming its output field differently would silently mis-normalize. + +Instead, every attack technique's ``run()`` method returns +``list[AttackResult]``: an explicit, frozen (immutable) Pydantic v2 model. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class Evaluation(BaseModel): + """A single evaluation/judgement attached to an :class:`AttackResult`.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str = "" + score: Optional[float] = None + success: Optional[bool] = None + notes: str = "" + metadata: Dict[str, Any] = Field(default_factory=dict) + + +def _evaluation_to_row(evaluation: Evaluation) -> Dict[str, Any]: + """Convert an :class:`Evaluation` back into a plain dict. + + If the evaluation only carries opaque ``metadata`` (all other fields at + their defaults) -- as produced by :meth:`AttackResult.from_row` for + legacy/technique-specific evaluation shapes that don't match the + ``Evaluation`` schema (e.g. ``{"classification": "SUCCESS"}``) -- the + original metadata dict is returned as-is so round-tripping through + :meth:`AttackResult.to_row` doesn't nest the original keys under a + ``"metadata"`` key. + """ + if ( + evaluation.name == "" + and evaluation.score is None + and evaluation.success is None + and evaluation.notes == "" + and evaluation.metadata + ): + return dict(evaluation.metadata) + return evaluation.model_dump() + + +class AttackResult(BaseModel): + """Typed, immutable representation of a single attack technique output row. + + Every attack technique returns ``list[AttackResult]`` from ``run()`` + instead of ad-hoc dicts/DataFrames/objects, so downstream orchestration + code no longer has to guess field names. + """ + + model_config = ConfigDict(frozen=True) + + goal: str = "" + prompt: str = "" + response: str = "" + evaluations: List[Evaluation] = Field(default_factory=list) + metadata: Dict[str, Any] = Field(default_factory=dict) + + @classmethod + def from_row(cls, row: Any) -> "AttackResult": + """Build an :class:`AttackResult` from a legacy heterogeneous row. + + Accepts a dict-like row (as produced by the existing pipeline steps) + and extracts the well-known fields, preserving everything else + (including the original raw values) in ``metadata`` so no + information is lost when converting back with :meth:`to_row`. + """ + if isinstance(row, AttackResult): + return row + if not isinstance(row, dict): + # Unknown/legacy row shape: preserve it as opaque metadata + # rather than raising, so callers get a valid (if minimal) + # AttackResult instead of silently dropping the row. + return cls(metadata={"_raw": row}) + + goal = row.get("goal") or "" + prompt = row.get("prompt") or row.get("prefix") or "" + response = row.get("response") or row.get("completion") or "" + + evaluations: List[Evaluation] = [] + raw_evaluations = row.get("evaluations") + if isinstance(raw_evaluations, list): + for item in raw_evaluations: + if isinstance(item, Evaluation): + evaluations.append(item) + elif isinstance(item, dict): + try: + evaluations.append(Evaluation(**item)) + except (TypeError, ValueError): + # Legacy/technique-specific evaluation shape (e.g. + # fields like "classification") that doesn't match + # the Evaluation schema: preserve it verbatim. + evaluations.append(Evaluation(metadata=dict(item))) + + return cls( + goal=goal, + prompt=prompt, + response=response, + evaluations=evaluations, + metadata=dict(row), + ) + + def to_row(self) -> Dict[str, Any]: + """Convert back to a plain ``dict`` for legacy dict-based code paths.""" + row = dict(self.metadata) + row["goal"] = self.goal + row["prompt"] = self.prompt + row["response"] = self.response + if self.evaluations: + row["evaluations"] = [_evaluation_to_row(e) for e in self.evaluations] + return row + + +def _extract_rows(results: Any) -> List[Any]: + """Extract a raw row list from a technique's ``run()`` return value. + + Shared by :func:`rows_to_attack_results` and :func:`flatten_run_result`. + Accepts ``None``, a list of rows, or a dict wrapping rows under one of + ``evaluated``/``rows``/``results``/``data``/``items`` (legacy whole-batch + shape). Returns rows unchanged (no per-row conversion). + """ + if results is None: + return [] + if isinstance(results, list): + return results + if isinstance(results, dict): + evaluated = results.get("evaluated") + if isinstance(evaluated, list): + return evaluated + for key in ("rows", "results", "data", "items"): + value = results.get(key) + if isinstance(value, list): + return value + return [] + return [] + + +def rows_to_attack_results(results: Any) -> List[AttackResult]: + """Normalize a technique's raw return value into ``list[AttackResult]``. + + This is the typed replacement for the old ``_normalize_attack_results`` + duck-typing helper. Accepts: + + - ``None`` -> ``[]`` + - a list of rows (dicts or :class:`AttackResult`) -> converted list + - a dict with an ``"evaluated"`` key (legacy baseline/static_template + shape) -> the ``"evaluated"`` rows, converted + - a dict with any of ``rows``/``results``/``data``/``items`` keys -> + those rows, converted + """ + return [AttackResult.from_row(r) for r in _extract_rows(results)] + + +def attack_results_to_rows(results: List[Any]) -> List[Any]: + """Convert ``list[AttackResult]`` back into plain dict rows. + + Used at the boundary with legacy dict-based downstream code (e.g. the + evaluator pipeline) that has not yet been migrated to the typed model. + Any item that isn't an :class:`AttackResult` (e.g. a legacy technique + already returning bare dicts/strings) is passed through unchanged. + """ + return [r.to_row() if isinstance(r, AttackResult) else r for r in results] + + +def flatten_run_result(results: Any) -> List[Any]: + """Flatten a technique's raw ``run()`` output into a list of rows. + + Unlike :func:`rows_to_attack_results`, this does **not** force every row + into an :class:`AttackResult` — it only extracts the row list from + legacy whole-batch dict shapes (``{"evaluated": [...], "summary": [...]}`` + etc.), preserving each row's original type. This is used internally by + the orchestrator when aggregating per-batch/per-goal ``run()`` calls, + where individual rows may already be :class:`AttackResult` instances or + (for legacy/third-party techniques) plain dicts/strings. + """ + return _extract_rows(results) diff --git a/hackagent/router/tracking/tracker.py b/hackagent/router/tracking/tracker.py index c05d590b..80a6e244 100644 --- a/hackagent/router/tracking/tracker.py +++ b/hackagent/router/tracking/tracker.py @@ -495,7 +495,6 @@ def _add_trace( content=sanitized_content, elapsed_s=trace_record["elapsed_s"], ) - self._record_failure(f"Goal {ctx.goal_index}: create trace", e) # Send to backend if enabled and we have a result_id if not self.is_enabled or not ctx.result_id: @@ -523,6 +522,7 @@ def _add_trace( f"Exception creating trace for goal {ctx.goal_index}: {e}", exc_info=True, ) + self._record_failure(f"Goal {ctx.goal_index}: create trace", e) return None diff --git a/tests/unit/attacks/advprefix/__init__.py b/tests/unit/attacks/advprefix/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/attacks/advprefix/test_attack_return_type.py b/tests/unit/attacks/advprefix/test_attack_return_type.py new file mode 100644 index 00000000..dd9a5a37 --- /dev/null +++ b/tests/unit/attacks/advprefix/test_attack_return_type.py @@ -0,0 +1,51 @@ +# Copyright 2026 - AI4I. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit test asserting AdvPrefixAttack.run() returns List[AttackResult].""" + +import unittest +from unittest.mock import MagicMock, patch + +from hackagent.attacks.techniques.advprefix.attack import AdvPrefixAttack +from hackagent.attacks.types import AttackResult + + +class TestAdvPrefixAttackReturnType(unittest.TestCase): + def test_run_empty_goals_returns_empty_list(self): + attack = AdvPrefixAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + self.assertEqual(attack.run([]), []) + + def test_run_returns_list_of_attack_result(self): + attack = AdvPrefixAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + coordinator = MagicMock() + coordinator.has_goal_tracking = False + coordinator.goal_tracker = None + + generation_output = [{"goal": "g1", "prompt": "p1"}] + final_output = [{"goal": "g1", "prompt": "p1", "response": "r1"}] + + with ( + patch.object(attack, "_initialize_coordinator", return_value=coordinator), + patch.object( + attack, + "_execute_pipeline", + side_effect=[generation_output, final_output], + ), + ): + results = attack.run(["g1"]) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g1") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/attacks/autodan_turbo/test_attack.py b/tests/unit/attacks/autodan_turbo/test_attack.py index 6f0a55e1..73fe372a 100644 --- a/tests/unit/attacks/autodan_turbo/test_attack.py +++ b/tests/unit/attacks/autodan_turbo/test_attack.py @@ -134,6 +134,10 @@ def _init_coord(*_args, **_kwargs): out = attack.run(["g"]) self.assertEqual(len(out), 1) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(out[0], AttackResult) + self.assertEqual(out[0].goal, "g") self.assertEqual(mock_warm_up.call_count, 1) self.assertEqual(mock_lifelong.call_count, 1) self.assertEqual(mock_eval.call_count, 1) diff --git a/tests/unit/attacks/baseline/__init__.py b/tests/unit/attacks/baseline/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/attacks/baseline/test_attack_return_type.py b/tests/unit/attacks/baseline/test_attack_return_type.py new file mode 100644 index 00000000..4a2298d3 --- /dev/null +++ b/tests/unit/attacks/baseline/test_attack_return_type.py @@ -0,0 +1,47 @@ +# Copyright 2026 - AI4I. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit test asserting BaselineAttack.run() returns List[AttackResult].""" + +import unittest +from unittest.mock import MagicMock, patch + +from hackagent.attacks.techniques.baseline.attack import BaselineAttack +from hackagent.attacks.types import AttackResult + + +class TestBaselineAttackReturnType(unittest.TestCase): + def test_run_empty_goals_returns_empty_list(self): + attack = BaselineAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + self.assertEqual(attack.run([]), []) + + def test_run_returns_list_of_attack_result(self): + attack = BaselineAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + coordinator = MagicMock() + coordinator.goal_tracker = None + + with ( + patch.object(attack, "_initialize_coordinator", return_value=coordinator), + patch.object( + attack, + "_execute_pipeline", + return_value=[{"goal": "g1", "response": "r1"}], + ), + ): + results = attack.run(["g1"]) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g1") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/attacks/bon/test_attack.py b/tests/unit/attacks/bon/test_attack.py index 7e5b0749..72a01777 100644 --- a/tests/unit/attacks/bon/test_attack.py +++ b/tests/unit/attacks/bon/test_attack.py @@ -184,6 +184,10 @@ def _init_coord(*_args, **_kwargs): mock_gen.assert_called_once() mock_eval.assert_called_once() self.assertEqual(len(results), 1) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "test") if __name__ == "__main__": diff --git a/tests/unit/attacks/cipherchat/test_attack.py b/tests/unit/attacks/cipherchat/test_attack.py index a2282f95..ef68d73e 100644 --- a/tests/unit/attacks/cipherchat/test_attack.py +++ b/tests/unit/attacks/cipherchat/test_attack.py @@ -214,6 +214,10 @@ def _init_coord(*_args, **_kwargs): mock_gen.assert_called_once() mock_eval.assert_called_once() self.assertEqual(len(results), 1) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "test") @patch("hackagent.attacks.techniques.cipherchat.attack.generation.execute") def test_run_no_generation_output(self, mock_gen): diff --git a/tests/unit/attacks/fc/test_attack.py b/tests/unit/attacks/fc/test_attack.py index 886cd896..f7bfdaa8 100644 --- a/tests/unit/attacks/fc/test_attack.py +++ b/tests/unit/attacks/fc/test_attack.py @@ -4,13 +4,14 @@ """Unit tests for FCAttack and tFCAttack classes.""" import unittest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from hackagent.attacks.techniques.fc.attack import ( FCAttack, tFCAttack, _recursive_update, ) +from hackagent.attacks.types import AttackResult class TestRecursiveUpdate(unittest.TestCase): @@ -85,6 +86,26 @@ def test_run_empty_goals(self): ) self.assertEqual(attack.run([]), []) + def test_run_returns_list_of_attack_result(self): + attack = FCAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + with ( + patch.object(attack, "_initialize_coordinator", return_value=MagicMock()), + patch.object( + attack, + "_execute_pipeline", + return_value=[{"goal": "g1", "response": "r1"}], + ), + ): + results = attack.run(["g1"]) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g1") + def test_get_pipeline_steps_returns_two_stages(self): attack = FCAttack( config={"output_dir": "./logs/runs"}, @@ -152,6 +173,26 @@ def test_run_empty_goals(self): ) self.assertEqual(attack.run([]), []) + def test_run_returns_list_of_attack_result(self): + attack = tFCAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + with ( + patch.object(attack, "_initialize_coordinator", return_value=MagicMock()), + patch.object( + attack, + "_execute_pipeline", + return_value=[{"goal": "g1", "response": "r1"}], + ), + ): + results = attack.run(["g1"]) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g1") + def test_get_pipeline_steps_returns_two_stages(self): attack = tFCAttack( config={"output_dir": "./logs/runs"}, diff --git a/tests/unit/attacks/flipattack/test_attack.py b/tests/unit/attacks/flipattack/test_attack.py index 16b6bb44..b6769dfd 100644 --- a/tests/unit/attacks/flipattack/test_attack.py +++ b/tests/unit/attacks/flipattack/test_attack.py @@ -118,6 +118,10 @@ def _init_coord(*_args, **_kwargs): self.assertEqual(len(out), 1) mock_generation.assert_called_once() mock_evaluation.assert_called_once() + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(out[0], AttackResult) + self.assertEqual(out[0].goal, "g1") if __name__ == "__main__": diff --git a/tests/unit/attacks/h4rm3l/test_attack.py b/tests/unit/attacks/h4rm3l/test_attack.py index 0635dde1..824b00e4 100644 --- a/tests/unit/attacks/h4rm3l/test_attack.py +++ b/tests/unit/attacks/h4rm3l/test_attack.py @@ -195,6 +195,10 @@ def _init_coord(*_args, **_kwargs): mock_gen.assert_called_once() mock_eval.assert_called_once() self.assertEqual(len(results), 1) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "test") if __name__ == "__main__": diff --git a/tests/unit/attacks/indirect_prompt_injection/__init__.py b/tests/unit/attacks/indirect_prompt_injection/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/attacks/indirect_prompt_injection/test_attack_return_type.py b/tests/unit/attacks/indirect_prompt_injection/test_attack_return_type.py new file mode 100644 index 00000000..6eead0b2 --- /dev/null +++ b/tests/unit/attacks/indirect_prompt_injection/test_attack_return_type.py @@ -0,0 +1,88 @@ +# Copyright 2026 - AI4I. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit test asserting IndirectPromptInjectionAttack.run() returns List[AttackResult].""" + +import unittest +from unittest.mock import MagicMock, patch + +from hackagent.attacks.techniques.indirect_prompt_injection.attack import ( + IndirectPromptInjectionAttack, +) +from hackagent.attacks.types import AttackResult + + +def _fake_create_router(backend, config, logger, router_name): + return MagicMock(), f"{router_name}_key" + + +class TestIndirectPromptInjectionAttackReturnType(unittest.TestCase): + def _make_attack(self): + with patch( + "hackagent.attacks.techniques.indirect_prompt_injection.attack.create_router", + side_effect=_fake_create_router, + ): + attack = IndirectPromptInjectionAttack( + config={"output_dir": "./logs/runs"}, + client=MagicMock(), + agent_router=MagicMock(), + ) + return attack + + def test_requires_client(self): + with self.assertRaises(ValueError): + IndirectPromptInjectionAttack( + config={}, client=None, agent_router=MagicMock() + ) + + def test_requires_agent_router(self): + with self.assertRaises(ValueError): + IndirectPromptInjectionAttack( + config={}, client=MagicMock(), agent_router=None + ) + + def test_run_returns_list_of_attack_result(self): + attack = self._make_attack() + + coordinator = MagicMock() + coordinator.has_goal_tracking = False + coordinator.goal_tracker = None + + with ( + patch.object(attack, "_initialize_coordinator", return_value=coordinator), + patch( + "hackagent.attacks.techniques.indirect_prompt_injection.attack.parse_documents", + return_value=[{"source": "doc1", "text": "some document text"}], + ), + patch.object( + attack, + "_run_single_goal", + return_value={ + "goal": "g1", + "evaluations": [{"classification": "SUCCESS"}], + }, + ), + ): + results = attack.run(["g1"]) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g1") + + def test_run_no_goals_raises(self): + attack = self._make_attack() + with self.assertRaises(ValueError): + attack.run([]) + + def test_run_no_documents_raises(self): + attack = self._make_attack() + with patch( + "hackagent.attacks.techniques.indirect_prompt_injection.attack.parse_documents", + return_value=[], + ): + with self.assertRaises(ValueError): + attack.run(["g1"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/attacks/mml/test_attack.py b/tests/unit/attacks/mml/test_attack.py index cc64212f..ff34a119 100644 --- a/tests/unit/attacks/mml/test_attack.py +++ b/tests/unit/attacks/mml/test_attack.py @@ -587,7 +587,12 @@ def test_run_calls_pipeline(self, mock_execute, mock_coordinator, mock_base_init mock_coordinator.assert_called_once() mock_execute.assert_called_once() - assert results == [{"goal": "test", "success": True}] + from hackagent.attacks.types import AttackResult + + assert len(results) == 1 + assert isinstance(results[0], AttackResult) + assert results[0].goal == "test" + assert results[0].metadata == {"goal": "test", "success": True} @patch("hackagent.attacks.techniques.base.BaseAttack.__init__", return_value=None) @patch("hackagent.attacks.techniques.mml.attack.MMLAttack._initialize_coordinator") diff --git a/tests/unit/attacks/pair/test_attack.py b/tests/unit/attacks/pair/test_attack.py index 61116112..994a96fd 100644 --- a/tests/unit/attacks/pair/test_attack.py +++ b/tests/unit/attacks/pair/test_attack.py @@ -213,6 +213,10 @@ def __call__(self, *_args, **_kwargs): self.assertEqual(len(results), 1) fake_coordinator.get_goal_context.assert_called_once_with(5) self.assertEqual(run_goal_mock.call_args.kwargs["goal_index"], 5) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "g") def test_single_goal_stops_immediately_on_jailbreak_score(self): dummy_attacker = MagicMock() diff --git a/tests/unit/attacks/pap/test_attack.py b/tests/unit/attacks/pap/test_attack.py index a5ba8510..277232c5 100644 --- a/tests/unit/attacks/pap/test_attack.py +++ b/tests/unit/attacks/pap/test_attack.py @@ -159,6 +159,10 @@ def _init_coord(*_args, **_kwargs): mock_gen.assert_called_once() mock_eval.assert_called_once() self.assertEqual(len(results), 1) + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(results[0], AttackResult) + self.assertEqual(results[0].goal, "test goal") if __name__ == "__main__": diff --git a/tests/unit/attacks/rag/test_attack.py b/tests/unit/attacks/rag/test_attack.py index 8f12de21..dc15dde2 100644 --- a/tests/unit/attacks/rag/test_attack.py +++ b/tests/unit/attacks/rag/test_attack.py @@ -276,18 +276,18 @@ def _run_with_strategy(self, strategy, extra_params=None): def test_run_inline_context_override(self): results = self._run_with_strategy("inline_context_override") self.assertEqual(len(results), 1) - self.assertIn("metrics", results[0]) - self.assertIn("asr", results[0]["metrics"]) + self.assertIn("metrics", results[0].metadata) + self.assertIn("asr", results[0].metadata["metrics"]) def test_run_append_hidden_directive(self): results = self._run_with_strategy("append_hidden_directive") self.assertEqual(len(results), 1) - self.assertGreaterEqual(results[0]["documents_poisoned"], 1) + self.assertGreaterEqual(results[0].metadata["documents_poisoned"], 1) def test_run_maximize_retrieval(self): results = self._run_with_strategy("maximize_retrieval") self.assertEqual(len(results), 1) - self.assertTrue(results[0]["evaluations"]) + self.assertTrue(results[0].evaluations) def test_run_manual_queries_vulnerable_mode(self): results = self._run_with_strategy( diff --git a/tests/unit/attacks/static_template/test_attack.py b/tests/unit/attacks/static_template/test_attack.py index 618eecb4..494d735f 100644 --- a/tests/unit/attacks/static_template/test_attack.py +++ b/tests/unit/attacks/static_template/test_attack.py @@ -55,7 +55,7 @@ def test_run_empty_goals(self): client=MagicMock(), agent_router=MagicMock(), ) - self.assertEqual(attack.run([]), {"evaluated": [], "summary": []}) + self.assertEqual(attack.run([]), []) @patch("hackagent.attacks.techniques.static_template.attack.evaluation.execute") @patch("hackagent.attacks.techniques.static_template.attack.generation.execute") @@ -83,8 +83,9 @@ def _init_coord(*_args, **_kwargs): with patch.object(attack, "_initialize_coordinator", side_effect=_init_coord): out = attack.run(["g1"]) - self.assertIn("evaluated", out) - self.assertIn("summary", out) + self.assertEqual(len(out), 1) + self.assertEqual(out[0].goal, "g1") + self.assertTrue(out[0].metadata.get("success")) mock_generation.assert_called_once() mock_evaluation.assert_called_once() diff --git a/tests/unit/attacks/tap/test_attack.py b/tests/unit/attacks/tap/test_attack.py index 4b3b506b..d9d321f2 100644 --- a/tests/unit/attacks/tap/test_attack.py +++ b/tests/unit/attacks/tap/test_attack.py @@ -95,6 +95,10 @@ def _init_coord(*_args, **_kwargs): self.assertEqual(len(out), 1) mock_generation.assert_called_once() mock_evaluation.assert_called_once() + from hackagent.attacks.types import AttackResult + + self.assertIsInstance(out[0], AttackResult) + self.assertEqual(out[0].goal, "g1") if __name__ == "__main__": diff --git a/tests/unit/attacks/test_error_propagation.py b/tests/unit/attacks/test_error_propagation.py index 259453cb..f4677d3f 100644 --- a/tests/unit/attacks/test_error_propagation.py +++ b/tests/unit/attacks/test_error_propagation.py @@ -19,7 +19,7 @@ # ============================================================================ -# 1. Orchestrator: _normalize_attack_results +# 1. attacks.types: rows_to_attack_results / attack_results_to_rows # ============================================================================ @@ -27,31 +27,72 @@ class TestNormalizeAttackResults(unittest.TestCase): """Test that dict-style results from static template are normalised to a list.""" def test_list_passthrough(self): - from hackagent.attacks.orchestrator import AttackOrchestrator + from hackagent.attacks.types import rows_to_attack_results data = [{"goal": "g1", "completion": "c1"}] - self.assertIs(AttackOrchestrator._normalize_attack_results(data), data) + result = rows_to_attack_results(data) + self.assertEqual(len(result), 1) + self.assertEqual(result[0].goal, "g1") + self.assertEqual(result[0].response, "c1") def test_dict_extracts_evaluated(self): - from hackagent.attacks.orchestrator import AttackOrchestrator + from hackagent.attacks.types import rows_to_attack_results evaluated = [{"goal": "g1", "completion": "c1"}] - result = AttackOrchestrator._normalize_attack_results( + result = rows_to_attack_results( {"evaluated": evaluated, "summary": [{"rate": 0.5}]} ) - self.assertIs(result, evaluated) + self.assertEqual(len(result), 1) + self.assertEqual(result[0].goal, "g1") def test_none_returns_empty(self): - from hackagent.attacks.orchestrator import AttackOrchestrator + from hackagent.attacks.types import rows_to_attack_results - self.assertEqual(AttackOrchestrator._normalize_attack_results(None), []) + self.assertEqual(rows_to_attack_results(None), []) def test_dict_without_evaluated_falls_back(self): - from hackagent.attacks.orchestrator import AttackOrchestrator + from hackagent.attacks.types import rows_to_attack_results rows = [{"goal": "g1"}] - result = AttackOrchestrator._normalize_attack_results({"rows": rows}) - self.assertEqual(result, rows) + result = rows_to_attack_results({"rows": rows}) + self.assertEqual(len(result), 1) + self.assertEqual(result[0].goal, "g1") + + def test_round_trip_to_row(self): + from hackagent.attacks.types import ( + attack_results_to_rows, + rows_to_attack_results, + ) + + rows = [{"goal": "g1", "completion": "c1", "extra": "x"}] + results = rows_to_attack_results(rows) + round_tripped = attack_results_to_rows(results) + self.assertEqual(round_tripped[0]["goal"], "g1") + self.assertEqual(round_tripped[0]["extra"], "x") + + def test_round_trip_opaque_evaluation_shape(self): + """Legacy evaluation shapes (e.g. RAG/indirect-prompt-injection's + ``{"classification": "SUCCESS"}``) that don't match the ``Evaluation`` + schema must round-trip back to their original flat shape rather than + being nested under a ``"metadata"`` key. + """ + from hackagent.attacks.types import ( + attack_results_to_rows, + rows_to_attack_results, + ) + + rows = [ + { + "goal": "g1", + "evaluations": [{"classification": "SUCCESS", "rationale": "ok"}], + } + ] + results = rows_to_attack_results(rows) + round_tripped = attack_results_to_rows(results) + self.assertEqual( + round_tripped[0]["evaluations"][0]["classification"], "SUCCESS" + ) + self.assertEqual(round_tripped[0]["evaluations"][0]["rationale"], "ok") # ============================================================================