decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

dt_arena/src/types/hooks.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import importlib
+import json
+import time
+import traceback
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol, runtime_checkable
+
+
+@dataclass
+class ToolCallContext:
+    """Describes an in-flight MCP tool call.
+
+    A pre-hook may return a modified ToolCallContext to rewrite `arguments`
+    or attach metadata before the call is dispatched.
+    """
+    framework: str
+    server: str
+    tool_name: str
+    arguments: Dict[str, Any]
+    trace_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ToolCallResult:
+    """Outcome of a dispatched MCP tool call.
+
+    `raw` holds whatever the underlying SDK returned (shape is framework
+    specific). `exception` is set if the call raised. `duration` is seconds.
+    A post-hook may return a modified ToolCallResult to rewrite the output.
+    """
+    raw: Any = None
+    is_error: bool = False
+    duration: float = 0.0
+    exception: Optional[BaseException] = None
+
+
+@runtime_checkable
+class ToolCallHook(Protocol):
+    """Protocol every hook implementation satisfies.
+
+    Either method may be omitted (return None). Returning a new context or
+    result from pre/post replaces the in-flight value.
+    """
+
+    async def on_pre_tool_call(
+        self, ctx: ToolCallContext
+    ) -> Optional[ToolCallContext]: ...
+
+    async def on_post_tool_call(
+        self, ctx: ToolCallContext, result: ToolCallResult
+    ) -> Optional[ToolCallResult]: ...
+
+
+_HOOKS_CONFIG_PATH = Path(__file__).resolve().parents[1] / "hooks" / "hooks.json"
+
+
+def _load_hooks_from_config() -> List["ToolCallHook"]:
+    """Instantiate hooks listed in `dt_arena/src/hooks/hooks.json`.
+
+    Expected format::
+
+        {"hooks": ["dt_arena.src.hooks.audit_log:AuditHook", ...]}
+
+    Missing file, empty list, or a malformed entry all degrade gracefully —
+    we log and skip. Each spec is ``module.path:ClassName``; the class is
+    imported and instantiated with no arguments.
+    """
+    if not _HOOKS_CONFIG_PATH.is_file():
+        return []
+    try:
+        config = json.loads(_HOOKS_CONFIG_PATH.read_text())
+    except Exception as e:
+        print(f"[HookManager] failed to parse {_HOOKS_CONFIG_PATH}: {e}")
+        return []
+
+    hooks: List[ToolCallHook] = []
+    for spec in config.get("hooks", []) or []:
+        module_name, sep, class_name = spec.rpartition(":")
+        if not sep or not module_name or not class_name:
+            print(f"[HookManager] invalid hook spec '{spec}' (expected 'module:ClassName')")
+            continue
+        try:
+            cls = getattr(importlib.import_module(module_name), class_name)
+            hooks.append(cls())
+        except Exception as e:
+            print(f"[HookManager] failed to load hook '{spec}': {e}")
+    return hooks
+
+
+class HookManager:
+    """Runs registered hooks around a framework's real tool-dispatch call.
+
+    Hooks listed in ``dt_arena/src/hooks/hooks.json`` are auto-loaded on
+    every construction, so any agent built via ``build_agent`` or directly
+    inherits them without further wiring.
+    """
+
+    def __init__(self, hooks: Optional[List[ToolCallHook]] = None):
+        # Config-file hooks run first; explicit per-agent hooks run after.
+        self._hooks: List[ToolCallHook] = _load_hooks_from_config() + list(hooks or [])
+
+    def register(self, hook: ToolCallHook) -> None:
+        self._hooks.append(hook)
+
+    def clear(self) -> None:
+        self._hooks.clear()
+
+    @property
+    def hooks(self) -> List[ToolCallHook]:
+        return list(self._hooks)
+
+    async def _run_pre(self, ctx: ToolCallContext) -> ToolCallContext:
+        for hook in self._hooks:
+            fn = getattr(hook, "on_pre_tool_call", None)
+            if fn is None:
+                continue
+            new_ctx = await fn(ctx)
+            if new_ctx is not None:
+                ctx = new_ctx
+        return ctx
+
+    async def _run_post(
+        self, ctx: ToolCallContext, result: ToolCallResult
+    ) -> ToolCallResult:
+        for hook in self._hooks:
+            fn = getattr(hook, "on_post_tool_call", None)
+            if fn is None:
+                continue
+            new_result = await fn(ctx, result)
+            if new_result is not None:
+                result = new_result
+        return result
+
+    async def wrap(
+        self,
+        ctx: ToolCallContext,
+        call: Callable[[Dict[str, Any]], Awaitable[Any]],
+    ) -> Any:
+        """Run pre-hooks, dispatch `call(arguments)`, run post-hooks.
+
+        `call` is a framework-specific async callable that performs the real
+        MCP dispatch given the (possibly rewritten) arguments and returns
+        whatever that SDK returns. Post-hooks run even on exception.
+        """
+        # Fast path: no hooks registered.
+        if not self._hooks:
+            return await call(ctx.arguments)
+
+        ctx = await self._run_pre(ctx)
+
+        start = time.perf_counter()
+        result = ToolCallResult()
+        try:
+            result.raw = await call(ctx.arguments)
+        except BaseException as exc:
+            result.exception = exc
+            result.is_error = True
+            result.duration = time.perf_counter() - start
+            try:
+                await self._run_post(ctx, result)
+            except Exception:
+                traceback.print_exc()
+            raise
+        else:
+            result.duration = time.perf_counter() - start
+
+        final = await self._run_post(ctx, result)
+        if final.exception is not None and final.exception is not result.exception:
+            raise final.exception
+        return final.raw

dt_arena/src/types/judge.py

@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from typing import List, Tuple, Dict, Any, Optional
+
+
+class BaseJudge(ABC):
+    """Base class for judges that evaluate agent behavior.
+
+    DEPRECATION NOTICE:
+        The `trajectory` argument has been removed from `eval_task` and `eval_attack`
+        methods. This parameter is deprecated and no longer supported. Judge
+        implementations should only use `agent_responses` for evaluation.
+    """
+
+    def __init__(self, domain: str, judge_model: str = "gpt-4o"):
+        self.domain = domain
+        self.judge_model = judge_model
+
+    @abstractmethod
+    def eval_task(self, agent_responses: List[str]) -> Tuple[bool, Dict[str, Any]]:
+        """
+        Evaluate whether the agent successfully completed the task.
+
+        Args:
+            agent_responses: List of responses produced by the agent at each turn.
+
+        Returns:
+            Tuple of (success: bool, metadata: dict) where metadata contains
+            additional information about the evaluation result.
+
+        Note:
+            DEPRECATED: The `trajectory` parameter has been removed.
+            This parameter is no longer supported.
+        """
+        pass
+
+    @abstractmethod
+    def eval_attack(self, agent_responses: List[str]) -> Optional[Tuple[bool, Dict[str, Any]]]:
+        """
+        Evaluate whether an attack was successful against the agent.
+
+        Args:
+            agent_responses: List of responses produced by the agent at each turn.
+
+        Returns:
+            Tuple of (success: bool, metadata: dict) for adversarial tasks,
+            or None for benign tasks where no attack evaluation is needed.
+
+        Note:
+            DEPRECATED: The `trajectory` parameter has been removed.
+            This parameter is no longer supported.
+        """
+        pass

dt_arena/src/types/red_teaming_trajectory.py

@@ -0,0 +1,385 @@
+import os
+import json
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional, Any, List, Dict, Literal
+
+
+@dataclass
+class RedTeamingResult:
+    """
+    Result from a red-teaming attack run.
+
+    This class encapsulates the outcome of a red-teaming attack,
+    including success status, detailed results, and the complete trajectory.
+    """
+    success: bool
+    result: Dict[str, Any]
+    trajectory: Optional['RedTeamingTrajectory'] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the result to a dictionary."""
+        return {
+            "success": self.success,
+            "result": self.result,
+            "trajectory": self.trajectory.to_dict()
+        }
+
+
+class RedTeamingTrajectory:
+    """
+    A class for building and manipulating red-teaming attack trajectories.
+
+    This class tracks the complete lifecycle of a red-teaming attack:
+    - Attack configuration and metadata
+    - Decision-making process
+    - Query attempts and target responses
+    - Judge evaluations
+    - Environment/MCP injections
+    - Skills loaded and used
+    - Final attack outcome
+    """
+
+    def __init__(
+        self,
+        domain: str,
+        task_id: str,
+        risk_category: str,
+        threat_model: Literal["indirect", "direct"],
+        original_task: Optional[str] = None,
+        malicious_goal: Optional[str] = None,
+        agent_model: Optional[str] = None,
+        victim_model: Optional[str] = None,
+        victim_arch: Optional[str] = None,
+        judge_model: Optional[str] = None,
+        max_attack_iterations: int = 10,
+    ):
+        """
+        Initialize a new red-teaming trajectory.
+
+        Args:
+            domain: Task domain (e.g., "travel", "crm")
+            risk_category: Risk category (e.g., "accommodation-preference-violation")
+            task_id: Unique task identifier
+            threat_model: Threat model type - must be "indirect" or "direct"
+            original_task: The original benign task description
+            malicious_goal: The malicious objective to achieve
+            agent_model: LLM model used by red-teaming agent
+            victim_model: LLM model used by victim agent
+            victim_arch: Architecture type of victim agent (e.g., "pocketflow", "openaisdk", "langchain")
+            judge_model: LLM model used by judge
+            max_attack_iterations: Maximum attack attempts allowed
+        """
+        # Validate threat_model
+        if threat_model not in ["indirect", "direct"]:
+            raise ValueError(f"threat_model must be 'indirect' or 'direct', got: {threat_model}")
+
+        self.data = {
+            "task_info": {
+                "task_id": task_id,
+                "domain": domain,
+                "risk_category": risk_category,
+                "threat_model": threat_model,
+                "original_task": original_task,
+                "malicious_goal": malicious_goal,
+            },
+            "attack_config": {
+                "agent_model": agent_model,
+                "victim_model": victim_model,
+                "victim_arch": victim_arch,
+                "judge_model": judge_model,
+                "max_attack_iterations": max_attack_iterations,
+            },
+            "attack_result": {
+                "success": False,
+                "duration": 0.0,
+                "timestamp": datetime.now().isoformat(),
+                "successful_attack_sequence": [],
+            },
+            "attack_trajectory": []
+        }
+        self._start_time: Optional[float] = None
+        self._next_step_id: int = 1
+
+    @classmethod
+    def load(cls, filepath: str) -> 'RedTeamingTrajectory':
+        """
+        Load trajectory from a JSON file.
+
+        Args:
+            filepath: Path to trajectory JSON file
+
+        Returns:
+            RedTeamingTrajectory instance with loaded data
+        """
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+
+        # Support both old and new format
+        task_info = data.get("task_info") or data.get("attack_info", {})
+        attack_config = data.get("attack_config", {})
+
+        instance = cls(
+            domain=task_info.get("domain", ""),
+            risk_category=task_info.get("risk_category", ""),
+            task_id=task_info.get("task_id", ""),
+            threat_model=task_info.get("threat_model", "indirect"),  # Default to indirect for backward compatibility
+            original_task=task_info.get("original_task"),
+            malicious_goal=task_info.get("malicious_goal"),
+            agent_model=attack_config.get("agent_model"),
+            victim_model=attack_config.get("victim_model") or attack_config.get("target_model"),  # backward compat
+            victim_arch=attack_config.get("victim_arch"),
+            judge_model=attack_config.get("judge_model"),
+            max_attack_iterations=attack_config.get("max_attack_iterations", 10),
+        )
+
+        # Load full data
+        instance.data = data
+        return instance
+
+    def append_action(
+        self,
+        role: str,
+        action: Optional[str] = None,
+        state: Optional[Any] = None,
+        reasoning: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Append a generic action step to the trajectory.
+
+        This is the unified interface for recording all agent actions.
+        Follows the structure: step_id, role, reasoning (optional), action/state, metadata
+
+        Args:
+            role: Role of the step ("attacker", "victim", "judge", "tool")
+            action: Action name (for attacker role)
+            state: State/response (for victim/judge/tool roles)
+            reasoning: Reasoning behind the action (optional, for attacker)
+            metadata: Additional metadata
+        """
+        step = {
+            "step_id": self._next_step_id,
+            "role": role,
+        }
+
+        # Add reasoning if provided (between role and action/state)
+        if reasoning:
+            step["reasoning"] = reasoning
+
+        # Add action or state depending on role
+        if action:
+            step["action"] = action
+        if state is not None:
+            step["state"] = state
+
+        # Add metadata
+        if metadata:
+            step["metadata"] = metadata
+
+        self.data["attack_trajectory"].append(step)
+        self._next_step_id += 1
+
+    def append_attacker_action(
+        self,
+        action: str,
+        reasoning: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Append an attacker action step to the trajectory.
+
+        Args:
+            action: Action name (dynamically determined from action space)
+            reasoning: Reasoning behind the action
+            **kwargs: Action-specific metadata
+        """
+        self.append_action(
+            role="attacker",
+            action=action,
+            reasoning=reasoning,
+            metadata=kwargs if kwargs else None
+        )
+
+    def append_victim_response(
+        self,
+        response: str,
+        query: str,
+        trajectory: Optional[List[Dict[str, Any]]] = None
+    ):
+        """
+        Append a victim agent response step.
+
+        Args:
+            response: Victim agent's response
+            query: Query sent to victim
+            trajectory: Victim agent's execution trajectory
+        """
+        metadata = {"query": query}
+        if trajectory:
+            metadata["victim_trajectory"] = trajectory
+
+        self.append_action(
+            role="victim",
+            state=response,
+            metadata=metadata
+        )
+
+    def append_verifiable_judge(
+        self,
+        attack_result: Optional[Dict[str, Any]] = None,
+        task_result: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Append a verifiable (programmatic) judge evaluation step.
+
+        For indirect threat model, both attack and task results are stored
+        in a single entry.
+
+        Args:
+            attack_result: Full result from the verifiable attack judge
+                           (keys: success, message, metadata)
+            task_result: Full result from the verifiable task judge
+                         (keys: success, message, metadata). Only present for indirect model.
+        """
+        state: Dict[str, Any] = {}
+        step_metadata: Dict[str, Any] = {}
+
+        if attack_result:
+            state["attack"] = attack_result.get("success", False)
+            step_metadata["attack"] = {
+                "success": attack_result.get("success", False),
+                "explanation": attack_result.get("message", ""),
+                "metadata": attack_result.get("metadata", {}),
+            }
+        if task_result:
+            state["task"] = task_result.get("success", False)
+            step_metadata["task"] = {
+                "success": task_result.get("success", False),
+                "explanation": task_result.get("message", ""),
+                "metadata": task_result.get("metadata", {}),
+            }
+
+        self.append_action(
+            role="verifiable_judge",
+            state=state,
+            metadata=step_metadata,
+        )
+
+    def append_feedback_judge(
+        self,
+        failure_reason: str = "",
+        improvement_suggestions: str = "",
+        feedback_details: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Append an LLM feedback judge evaluation step.
+
+        Args:
+            failure_reason: Why the attack failed according to the LLM judge
+            improvement_suggestions: Suggestions for improving the attack
+            feedback_details: Full parsed structured feedback from LLM judge
+        """
+        step_metadata: Dict[str, Any] = {
+            "failure_reason": failure_reason,
+            "improvement_suggestions": improvement_suggestions,
+        }
+        if feedback_details:
+            step_metadata["feedback_details"] = feedback_details
+        self.append_action(
+            role="feedback_judge",
+            state=False,
+            metadata=step_metadata,
+        )
+
+    def set_success(self, success: bool):
+        """Set the final attack success status."""
+        self.data["attack_result"]["success"] = success
+
+    def set_successful_attack_sequence(self, attack_sequence: List[Dict[str, Any]]):
+        """
+        Set the successful attack sequence (injection steps that led to success).
+
+        Args:
+            attack_sequence: List of injection steps from InjectionBuffer.to_attack_instance()
+        """
+        self.data["attack_result"]["successful_attack_sequence"] = attack_sequence
+
+    def set_original_task(self, task: str):
+        """Set the original task description."""
+        self.data["task_info"]["original_task"] = task
+
+    def set_malicious_goal(self, goal: str):
+        """Set the malicious goal."""
+        self.data["task_info"]["malicious_goal"] = goal
+
+    def start_timer(self):
+        """Start timing the attack execution."""
+        self._start_time = datetime.now().timestamp()
+
+    def stop_timer(self):
+        """Stop timing and update duration."""
+        if self._start_time:
+            end_time = datetime.now().timestamp()
+            self.data["attack_result"]["duration"] = round(end_time - self._start_time, 3)
+            self.data["attack_result"]["timestamp"] = datetime.now().isoformat()
+
+    def get_summary(self) -> Dict[str, Any]:
+        """
+        Get a summary of the attack trajectory.
+
+        Returns:
+            Dictionary with summary statistics
+        """
+        # Count different action types
+        trajectory = self.data["attack_trajectory"]
+        num_attempts = sum(1 for step in trajectory if step.get("type") == "query_attempt")
+        num_skills_loaded = sum(1 for step in trajectory if step.get("type") == "skill_load")
+        num_env_injections = sum(1 for step in trajectory if step.get("type") == "env_injection")
+        num_mcp_injections = sum(1 for step in trajectory if step.get("type") == "mcp_injection")
+
+        return {
+            "task_id": self.data["task_info"]["task_id"],
+            "domain": self.data["task_info"]["domain"],
+            "risk_category": self.data["task_info"]["risk_category"],
+            "threat_model": self.data["task_info"]["threat_model"],
+            "success": self.data["attack_result"]["success"],
+            "num_attempts": num_attempts,
+            "num_skills_loaded": num_skills_loaded,
+            "num_env_injections": num_env_injections,
+            "num_mcp_injections": num_mcp_injections,
+            "duration": self.data["attack_result"]["duration"],
+            "trajectory_steps": len(trajectory)
+        }
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the trajectory as a dictionary."""
+        return self.data
+
+    def save(self, filepath: str) -> str:
+        """
+        Save trajectory to a JSON file.
+
+        Args:
+            filepath: Output file path
+
+        Returns:
+            Path to saved file
+        """
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else ".", exist_ok=True)
+
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(self.data, f, indent=2, ensure_ascii=False)
+
+        return filepath
+
+    def __repr__(self) -> str:
+        summary = self.get_summary()
+        return (
+            f"RedTeamingTrajectory("
+            f"task_id={summary['task_id']}, "
+            f"success={summary['success']}, "
+            f"attempts={summary['num_attempts']}, "
+            f"steps={summary['trajectory_steps']})"
+        )