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

utils/agent_helpers.py

@@ -0,0 +1,308 @@
+from typing import Any, Dict, List, Optional, Type
+
+
+OS_FILESYSTEM_CLAUDE_SDK_DISALLOWED_TOOLS = [
+    "Bash",
+    "Read",
+    "Write",
+    "Edit",
+    "MultiEdit",
+    "Glob",
+    "Grep",
+    "NotebookEdit",
+    "AskUserQuestion",
+]
+
+
+OS_FILESYSTEM_OPENCLAW_DISALLOWED_TOOLS = [
+    "group:fs",
+    "group:runtime",
+    "group:web",
+    "group:memory",
+    "group:ui",
+    "group:sessions",
+    "group:automation",
+    "group:messaging",
+    "agents_list",
+    "image",
+    "nodes",
+    "tts",
+]
+
+
+def get_default_disallowed_tools(agent_type: Optional[str], domain: Optional[str]) -> Optional[List[str]]:
+    """Return domain-specific native tools to disable when none are provided explicitly."""
+    if (domain or "").lower() != "os-filesystem":
+        return None
+
+    agent = (agent_type or "").lower()
+    if agent == "claudesdk":
+        return list(OS_FILESYSTEM_CLAUDE_SDK_DISALLOWED_TOOLS)
+    if agent == "openclaw":
+        return list(OS_FILESYSTEM_OPENCLAW_DISALLOWED_TOOLS)
+    return None
+
+_AgentConfig = None
+_RuntimeConfig = None
+_TaskConfig = None
+_AttackConfig = None
+_OpenAISDKAgent = None
+_LangChainAgent = None
+_ClaudeSDKAgent = None
+_GoogleADKAgent = None
+_MCPReactAgent = None
+_OpenClawAgent = None
+
+# Native agent wrappers (lazy imports)
+_OpenAISDKNativeWrapper = None
+_MCPReactAgentNativeWrapper = None
+
+
+def _ensure_imports():
+    """Lazy import of agent-related types."""
+    global _AgentConfig, _RuntimeConfig, _TaskConfig, _AttackConfig
+    global _OpenAISDKAgent, _LangChainAgent, _ClaudeSDKAgent, _GoogleADKAgent, _MCPReactAgent, _OpenClawAgent
+    global _OpenAISDKNativeWrapper, _MCPReactAgentNativeWrapper
+    
+    if _AgentConfig is None:
+        from dt_arena.src.types.agent import AgentConfig, RuntimeConfig
+        from dt_arena.src.types.task import AttackConfig, TaskConfig
+        _AgentConfig = AgentConfig
+        _RuntimeConfig = RuntimeConfig
+        _TaskConfig = TaskConfig
+        _AttackConfig = AttackConfig
+    
+    # Import OpenAI SDK agent
+    if _OpenAISDKAgent is None:
+        try:
+            from agent.openaisdk import OpenAISDKAgent
+            _OpenAISDKAgent = OpenAISDKAgent
+        except ImportError as e:
+            print(f"[WARNING] OpenAI SDK agent not available: {e}")
+    
+    # Import LangChain agent
+    if _LangChainAgent is None:
+        try:
+            from agent.langchain import LangChainAgent
+            _LangChainAgent = LangChainAgent
+        except ImportError as e:
+            print(f"[WARNING] LangChain agent not available: {e}")
+    
+    # Import Claude SDK agent
+    if _ClaudeSDKAgent is None:
+        try:
+            from agent.claudesdk import ClaudeSDKAgent
+            _ClaudeSDKAgent = ClaudeSDKAgent
+        except ImportError as e:
+            print(f"[WARNING] Claude SDK agent not available: {e}")
+
+    # Import Google ADK agent
+    if _GoogleADKAgent is None:
+        try:
+            from agent.googleadk import GoogleADKAgent
+            _GoogleADKAgent = GoogleADKAgent
+        except ImportError as e:
+            print(f"[WARNING] Google ADK agent not available: {e}")
+
+    # Import PocketFlow agent
+    if _MCPReactAgent is None:
+        try:
+            from agent.pocketflow import MCPReactAgent
+            _MCPReactAgent = MCPReactAgent
+        except ImportError as e:
+            print(f"[WARNING] PocketFlow agent not available: {e}")
+
+    # Import OpenClaw agent
+    if _OpenClawAgent is None:
+        try:
+            from agent.openclaw import OpenClawAgent
+            _OpenClawAgent = OpenClawAgent
+        except ImportError as e:
+            print(f"[WARNING] OpenClaw agent not available: {e}")
+
+    # Import native agent wrappers
+    if _OpenAISDKNativeWrapper is None:
+        try:
+            from agent.openaisdk.src.agent_wrapper import OpenAISDKNativeWrapper
+            _OpenAISDKNativeWrapper = OpenAISDKNativeWrapper
+        except ImportError as e:
+            print(f"[WARNING] OpenAI SDK native wrapper not available: {e}")
+
+    if _MCPReactAgentNativeWrapper is None:
+        try:
+            from agent.pocketflow.src.agent_wrapper import MCPReactAgentNativeWrapper
+            _MCPReactAgentNativeWrapper = MCPReactAgentNativeWrapper
+        except ImportError as e:
+            print(f"[WARNING] PocketFlow native wrapper not available: {e}")
+
+
+def _get_agent_registry() -> Dict[str, Type]:
+    """Get the agent registry with lazy imports."""
+    _ensure_imports()
+    registry = {}
+    
+    if _OpenAISDKAgent is not None:
+        registry["openaisdk"] = _OpenAISDKAgent
+    if _LangChainAgent is not None:
+        registry["langchain"] = _LangChainAgent
+    if _ClaudeSDKAgent is not None:
+        registry["claudesdk"] = _ClaudeSDKAgent
+    if _GoogleADKAgent is not None:
+        registry["googleadk"] = _GoogleADKAgent
+    if _MCPReactAgent is not None:
+        registry["pocketflow"] = _MCPReactAgent
+    if _OpenClawAgent is not None:
+        registry["openclaw"] = _OpenClawAgent
+
+    return registry
+
+
+# Agent registry - lazily populated
+AGENT_REGISTRY: Dict[str, Type] = {}
+
+
+def build_agent(
+    agent_type: Optional[str] = None,
+    agent_cfg: Optional[Any] = None,  # AgentConfig
+    runtime_cfg: Optional[Any] = None,  # RuntimeConfig
+    *,
+    native_agent: Optional[Any] = None,  # Pre-built native SDK agent
+) -> Any:
+    """Instantiate the requested agent implementation.
+
+    Args:
+        agent_type: Type of agent to build (e.g., "openaisdk", "langchain").
+            Required if native_agent is not provided.
+        agent_cfg: AgentConfig instance. For native agents, this provides
+            additional benchmark MCP servers to add.
+        runtime_cfg: RuntimeConfig instance with model, temperature, max_turns,
+            output_dir, and mcp_injection settings.
+        native_agent: Pre-built native SDK agent to wrap. When provided,
+            the framework is auto-detected and the agent is wrapped for
+            evaluation.
+
+    Returns:
+        Initialized agent instance
+
+    Raises:
+        ValueError: If neither agent_type nor native_agent is provided,
+            or if the native agent framework cannot be detected.
+    """
+    _ensure_imports()
+
+    # Use provided runtime_cfg or create default
+    if runtime_cfg is None:
+        runtime_cfg = _RuntimeConfig()
+
+    # Mode 2: Wrap a pre-built native agent
+    if native_agent is not None:
+        return _wrap_native_agent(
+            native_agent=native_agent,
+            agent_cfg=agent_cfg,
+            runtime_cfg=runtime_cfg,
+        )
+
+    # Mode 1: Build from scratch (existing behavior)
+    if agent_type is None:
+        raise ValueError(
+            "Either 'agent_type' or 'native_agent' must be provided. "
+            "Use agent_type to build from scratch, or native_agent to wrap "
+            "a pre-built agent."
+        )
+
+    registry = _get_agent_registry()
+    agent_type = agent_type.lower()
+
+    if agent_type not in registry:
+        raise ValueError(f"Unknown agent type '{agent_type}'. Available: {list(registry.keys())}")
+
+    AgentCls = registry[agent_type]
+
+    return AgentCls(
+        agent_config=agent_cfg,
+        runtime_config=runtime_cfg,
+    )
+
+
+def _wrap_native_agent(
+    native_agent: Any,
+    agent_cfg: Optional[Any],
+    runtime_cfg: Any,
+) -> Any:
+    """
+    Wrap a pre-built native SDK agent for evaluation.
+
+    Auto-detects the agent framework and returns the appropriate wrapper.
+
+    Args:
+        native_agent: Pre-built agent from a native SDK
+        agent_cfg: AgentConfig with benchmark MCP servers to add
+        runtime_cfg: RuntimeConfig with model settings, output paths, injections
+
+    Returns:
+        Wrapped agent instance
+
+    Raises:
+        ValueError: If framework cannot be detected or no wrapper is available
+    """
+    from utils.agent_wrapper import detect_native_framework
+
+    # Detect framework
+    framework = detect_native_framework(native_agent)
+
+    if framework is None:
+        raise ValueError(
+            f"Could not detect framework for agent of type {type(native_agent).__name__}. "
+            f"Supported native agent types: OpenAI SDK Agent, LangChain Agent, "
+            f"Claude SDK Client, Google ADK Agent, PocketFlow NativeMCPReactAgent."
+        )
+
+    # Get the appropriate wrapper
+    if framework == "openaisdk":
+        if _OpenAISDKNativeWrapper is None:
+            raise ValueError(
+                "OpenAI SDK native wrapper is not available. "
+                "Please ensure 'agents' package is installed."
+            )
+        return _OpenAISDKNativeWrapper(
+            native_agent=native_agent,
+            agent_config=agent_cfg,
+            runtime_config=runtime_cfg,
+        )
+    elif framework == "langchain":
+        # TODO: Implement LangChainNativeWrapper
+        raise NotImplementedError(
+            f"Native agent wrapper for '{framework}' is not yet implemented. "
+            f"Currently supported: openaisdk, pocketflow"
+        )
+    elif framework == "claudesdk":
+        # TODO: Implement ClaudeSDKNativeWrapper
+        raise NotImplementedError(
+            f"Native agent wrapper for '{framework}' is not yet implemented. "
+            f"Currently supported: openaisdk, pocketflow"
+        )
+    elif framework == "googleadk":
+        # TODO: Implement GoogleADKNativeWrapper
+        raise NotImplementedError(
+            f"Native agent wrapper for '{framework}' is not yet implemented. "
+            f"Currently supported: openaisdk, pocketflow"
+        )
+    elif framework == "pocketflow":
+        if _MCPReactAgentNativeWrapper is None:
+            raise ValueError(
+                "PocketFlow native wrapper is not available. "
+                "Please ensure 'pocketflow' package is installed."
+            )
+        return _MCPReactAgentNativeWrapper(
+            native_agent=native_agent,
+            agent_config=agent_cfg,
+            runtime_config=runtime_cfg,
+        )
+    else:
+        raise ValueError(f"No wrapper available for framework '{framework}'.")
+
+
+def get_agent_choices() -> list[str]:
+    """Get sorted list of available agent types."""
+    _ensure_imports()
+    return sorted(_get_agent_registry().keys())

utils/agent_wrapper.py

@@ -0,0 +1,189 @@
+"""
+Native Agent Detection Utilities
+
+This module provides helper functions for detecting the framework of
+pre-built native agents. The actual wrapper classes are located in
+their respective agent folders (e.g., agent/openaisdk/src/native_wrapper.py).
+
+Usage:
+    from utils.agent_wrapper import detect_native_framework
+
+    native_agent = ...  # User's pre-built agent
+    framework = detect_native_framework(native_agent)
+    # Returns: "openaisdk", "langchain", "claudesdk", "googleadk", or None
+"""
+
+from typing import Any, Optional
+
+
+def is_openai_sdk_native_agent(agent: Any) -> bool:
+    """
+    Detect if agent is a native OpenAI SDK Agent.
+
+    Args:
+        agent: Object to check
+
+    Returns:
+        True if this is an OpenAI Agents SDK agent
+    """
+    # Check for OpenAI Agents SDK (from 'agents' package)
+    try:
+        from agents import Agent as OpenAIAgent
+        if isinstance(agent, OpenAIAgent):
+            return True
+    except ImportError:
+        pass
+
+    # Duck-type check: OpenAI SDK agents have these specific attributes
+    return (
+        hasattr(agent, 'name') and
+        hasattr(agent, 'instructions') and
+        hasattr(agent, 'model') and
+        hasattr(agent, 'mcp_servers') and
+        hasattr(agent, 'model_settings') and
+        hasattr(agent, 'clone')
+    )
+
+
+def is_langchain_native_agent(agent: Any) -> bool:
+    """
+    Detect if agent is a native LangChain agent/graph.
+
+    Args:
+        agent: Object to check
+
+    Returns:
+        True if this is a LangChain agent
+    """
+    # Check for LangChain CompiledGraph (from create_agent)
+    try:
+        from langgraph.graph.state import CompiledGraph, CompiledStateGraph
+        if isinstance(agent, (CompiledGraph, CompiledStateGraph)):
+            return True
+    except ImportError:
+        pass
+
+    # Check for legacy AgentExecutor
+    try:
+        from langchain.agents import AgentExecutor
+        if isinstance(agent, AgentExecutor):
+            return True
+    except ImportError:
+        pass
+
+    # Duck-type check
+    return (
+        hasattr(agent, 'invoke') and
+        hasattr(agent, 'ainvoke') and
+        (hasattr(agent, 'get_graph') or hasattr(agent, 'nodes'))
+    )
+
+
+def is_claude_sdk_native_agent(agent: Any) -> bool:
+    """
+    Detect if agent is a native Claude SDK agent.
+
+    Args:
+        agent: Object to check
+
+    Returns:
+        True if this is a Claude SDK agent
+    """
+    try:
+        from claude_code_sdk import Client
+        if isinstance(agent, Client):
+            return True
+    except ImportError:
+        pass
+
+    # Duck-type check
+    return (
+        hasattr(agent, 'query') and
+        hasattr(agent, 'receive_response') and
+        hasattr(agent, '_model')
+    )
+
+
+def is_google_adk_native_agent(agent: Any) -> bool:
+    """
+    Detect if agent is a native Google ADK agent.
+
+    Args:
+        agent: Object to check
+
+    Returns:
+        True if this is a Google ADK agent
+    """
+    try:
+        from google.adk.agents import Agent as GoogleAgent, LlmAgent
+        if isinstance(agent, (GoogleAgent, LlmAgent)):
+            return True
+    except ImportError:
+        pass
+
+    # Duck-type check
+    return (
+        hasattr(agent, 'model') and
+        hasattr(agent, 'instruction') and
+        hasattr(agent, 'tools') and
+        hasattr(type(agent), '__module__') and
+        type(agent).__module__.startswith('google.adk')
+    )
+
+
+def is_pocketflow_native_agent(agent: Any) -> bool:
+    """
+    Detect if agent is a native PocketFlow agent (NativeMCPReactAgent).
+
+    Args:
+        agent: Object to check
+
+    Returns:
+        True if this is a PocketFlow NativeMCPReactAgent
+    """
+    # Type check for NativeMCPReactAgent
+    try:
+        from agent.pocketflow import NativeMCPReactAgent
+        if isinstance(agent, NativeMCPReactAgent):
+            return True
+    except ImportError:
+        pass
+
+    # Duck-type check: NativeMCPReactAgent has these specific attributes
+    return (
+        hasattr(agent, 'mcp_servers') and
+        hasattr(agent, 'tool_to_server') and
+        hasattr(agent, 'all_tools') and
+        isinstance(getattr(agent, 'mcp_servers', None), dict) and
+        isinstance(getattr(agent, 'tool_to_server', None), dict) and
+        isinstance(getattr(agent, 'all_tools', None), list)
+    )
+
+
+def detect_native_framework(agent: Any) -> Optional[str]:
+    """
+    Detect which framework a native agent belongs to.
+
+    Args:
+        agent: A pre-built native SDK agent instance
+
+    Returns:
+        Framework name ("openaisdk", "langchain", "claudesdk", "googleadk", "pocketflow")
+        or None if not detected
+    """
+    detectors = [
+        ("openaisdk", is_openai_sdk_native_agent),
+        ("langchain", is_langchain_native_agent),
+        ("claudesdk", is_claude_sdk_native_agent),
+        ("googleadk", is_google_adk_native_agent),
+        ("pocketflow", is_pocketflow_native_agent),
+    ]
+
+    for framework, detector in detectors:
+        try:
+            if detector(agent):
+                return framework
+        except Exception:
+            continue
+
+    return None

utils/compose_utils.py

@@ -0,0 +1,135 @@
+import os
+import subprocess
+from pathlib import Path
+from typing import List, Optional, Union
+
+
+def get_project_name(env_name: str) -> Optional[str]:
+    """
+    Get the Docker compose project name for a given environment.
+    
+    In pooled mode, the project name is passed via environment variable
+    like GMAIL_PROJECT_NAME, SLACK_PROJECT_NAME, etc.
+    
+    Args:
+        env_name: Environment name (gmail, slack, atlassian, etc.)
+    
+    Returns:
+        Project name if found, None otherwise
+    """
+    # Check for pool-provided project name (e.g., GMAIL_PROJECT_NAME)
+    env_var = f"{env_name.upper().replace('-', '_')}_PROJECT_NAME"
+    project_name = os.getenv(env_var, "")
+    if project_name:
+        return project_name
+    
+    # Fallback to base project name
+    base_project = os.getenv("COMPOSE_PROJECT_NAME", "")
+    if base_project:
+        return f"{base_project}-{env_name}"
+    
+    return None
+
+
+def get_project_args(env_name: str) -> List[str]:
+    """
+    Get docker compose project name arguments.
+    
+    Args:
+        env_name: Environment name
+    
+    Returns:
+        List like ["-p", "project_name"] or empty list
+    """
+    project_name = get_project_name(env_name)
+    if project_name:
+        return ["-p", project_name]
+    return []
+
+
+def run_compose(
+    env_name: str,
+    compose_file: Union[str, Path],
+    args: List[str],
+    cwd: Optional[Union[str, Path]] = None,
+    check: bool = True,
+    verbose: bool = True,
+) -> subprocess.CompletedProcess:
+    """
+    Run docker compose command with correct project name.
+    
+    Args:
+        env_name: Environment name (gmail, slack, etc.)
+        compose_file: Path to docker-compose.yml
+        args: Additional arguments for docker compose
+        cwd: Working directory
+        check: Raise exception on non-zero exit
+        verbose: Print command before running
+    
+    Returns:
+        CompletedProcess instance
+    """
+    project_args = get_project_args(env_name)
+    cmd = ["docker", "compose"] + project_args + ["-f", str(compose_file)] + args
+    
+    if verbose:
+        print("$ " + " ".join(cmd))
+    
+    return subprocess.run(cmd, cwd=cwd, check=check)
+
+
+def run_compose_exec(
+    env_name: str,
+    compose_file: Union[str, Path],
+    service: str,
+    command: List[str],
+    cwd: Optional[Union[str, Path]] = None,
+    check: bool = True,
+    verbose: bool = True,
+) -> subprocess.CompletedProcess:
+    """
+    Run command inside a docker compose service.
+    
+    Args:
+        env_name: Environment name (gmail, slack, etc.)
+        compose_file: Path to docker-compose.yml
+        service: Service name to exec into
+        command: Command to run
+        cwd: Working directory
+        check: Raise exception on non-zero exit
+        verbose: Print command before running
+    
+    Returns:
+        CompletedProcess instance
+    """
+    args = ["exec", "-T", service] + command
+    return run_compose(env_name, compose_file, args, cwd=cwd, check=check, verbose=verbose)
+
+
+def run_compose_cp(
+    env_name: str,
+    compose_file: Union[str, Path],
+    src: str,
+    dst: str,
+    cwd: Optional[Union[str, Path]] = None,
+    check: bool = True,
+    verbose: bool = True,
+) -> subprocess.CompletedProcess:
+    """
+    Copy files to/from a docker compose service.
+    
+    Args:
+        env_name: Environment name (gmail, slack, etc.)
+        compose_file: Path to docker-compose.yml
+        src: Source path (local or service:path)
+        dst: Destination path (local or service:path)
+        cwd: Working directory
+        check: Raise exception on non-zero exit
+        verbose: Print command before running
+    
+    Returns:
+        CompletedProcess instance
+    """
+    args = ["cp", src, dst]
+    return run_compose(env_name, compose_file, args, cwd=cwd, check=check, verbose=verbose)
+

utils/config.py

@@ -0,0 +1,77 @@
+import importlib.resources
+import os
+import sys
+from pathlib import Path
+
+
+def _find_project_root() -> Path:
+    # 1. Explicit override — set DTAP_ROOT when running from Docker / CI / installed wheel.
+    if env_root := os.environ.get("DTAP_ROOT"):
+        root = Path(env_root).resolve()
+        if str(root) not in sys.path:
+            sys.path.insert(0, str(root))
+        return root
+
+    # 2. Installed package: locate via the bundled dt_arena/config/env.yaml.
+    #    importlib.resources works whether the package is pip-installed or editable.
+    try:
+        ref = importlib.resources.files("dt_arena") / "config" / "env.yaml"
+        with importlib.resources.as_file(ref) as p:
+            # p = <prefix>/dt_arena/config/env.yaml  →  parents[2] = <prefix>
+            root = p.parents[2]
+            if str(root) not in sys.path:
+                sys.path.insert(0, str(root))
+            return root
+    except Exception:
+        pass
+
+    # 3. Source checkout: walk up until we find pyproject.toml.
+    for parent in Path(__file__).resolve().parents:
+        if (parent / "pyproject.toml").exists():
+            if str(parent) not in sys.path:
+                sys.path.insert(0, str(parent))
+            return parent
+
+    # 4. Last resort.
+    guess = Path(__file__).resolve().parents[1]
+    if str(guess) not in sys.path:
+        sys.path.insert(0, str(guess))
+    return guess
+
+
+PROJECT_ROOT = _find_project_root()
+
+# Configuration file paths (bundled as package-data in the wheel)
+ENV_CONFIG_PATH = PROJECT_ROOT / "dt_arena" / "config" / "env.yaml"
+MCP_CONFIG_PATH = PROJECT_ROOT / "dt_arena" / "config" / "mcp.yaml"
+INJECTION_MCP_CONFIG_PATH = PROJECT_ROOT / "dt_arena" / "config" / "injection_mcp.yaml"
+
+# Bundled benchmark task lists (benchmark/<domain>/{benign,direct,indirect}.jsonl)
+BENCHMARK_ROOT = PROJECT_ROOT / "benchmark"
+
+# Evaluation script path
+TASK_RUNNER_PATH = PROJECT_ROOT / "eval" / "task_runner.py"
+
+
+def resolve_benchmark_task_list(
+    domain: str | None = None,
+    task_type: str | None = None,
+    threat_model: str | None = None,
+) -> Path:
+    """Resolve the bundled benchmark JSONL/dir matching the given filters.
+
+    - No filters       → BENCHMARK_ROOT (all 14 domains × 3 task lists)
+    - domain only      → BENCHMARK_ROOT/<domain>/             (all 3 lists)
+    - +task_type=benign → BENCHMARK_ROOT/<domain>/benign.jsonl
+    - +threat_model    → BENCHMARK_ROOT/<domain>/<direct|indirect>.jsonl
+    """
+    if not domain:
+        return BENCHMARK_ROOT
+
+    domain_dir = BENCHMARK_ROOT / domain
+    if task_type == "benign":
+        return domain_dir / "benign.jsonl"
+    if threat_model in ("direct", "indirect"):
+        return domain_dir / f"{threat_model}.jsonl"
+    return domain_dir
+