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

cli/main.py

@@ -0,0 +1,182 @@
+"""Top-level Typer application wired as the `dtap` console script."""
+import sys
+from pathlib import Path
+
+import typer
+import yaml
+from rich.console import Console
+from rich.markup import escape
+from rich.table import Table
+
+from .scaffold import SUPPORTED_FRAMEWORKS, scaffold
+
+app = typer.Typer(
+    name="dtap",
+    help="DecodingTrust Agent Platform: run red-teaming evaluations against AI agents.",
+    add_completion=False,
+    no_args_is_help=True,
+)
+
+info_app = typer.Typer(
+    name="info",
+    help="Inspect supported domains and environments.",
+    add_completion=False,
+    no_args_is_help=True,
+)
+agent_app = typer.Typer(
+    name="agent",
+    help="List supported agent backends and scaffold new ones.",
+    add_completion=False,
+    no_args_is_help=True,
+)
+app.add_typer(info_app, name="info")
+app.add_typer(agent_app, name="agent")
+
+
+# Let `--help` flow through to the underlying argparse so users see real flags.
+_PASSTHROUGH = {
+    "allow_extra_args": True,
+    "ignore_unknown_options": True,
+    "help_option_names": [],
+}
+
+
+@app.command(
+    "eval",
+    context_settings=_PASSTHROUGH,
+    help="Run a parallel evaluation. Use 'dtap eval --help' for all flags.",
+)
+def cmd_eval(ctx: typer.Context) -> None:
+    from eval.evaluation import main
+
+    sys.argv = ["dtap eval"] + ctx.args
+    main()
+
+
+# ── `dtap info …` subcommands ────────────────────────────────────────────────
+
+@info_app.command("domain", help="List benchmark domains with task counts.")
+def info_domain() -> None:
+    from utils.config import BENCHMARK_ROOT
+
+    console = Console()
+    table = Table(title="Benchmark domains", show_lines=False)
+    table.add_column("Domain", style="bold cyan")
+    table.add_column("Benign", justify="right")
+    table.add_column("Direct", justify="right")
+    table.add_column("Indirect", justify="right")
+    table.add_column("Total", justify="right", style="bold")
+
+    def _count(p):
+        if not p.exists():
+            return 0
+        with p.open() as f:
+            return sum(1 for line in f if line.strip())
+
+    domains = sorted(
+        p.name for p in BENCHMARK_ROOT.iterdir()
+        if p.is_dir() and not p.name.startswith("_")
+    )
+    grand_total = 0
+    for d in domains:
+        b = _count(BENCHMARK_ROOT / d / "benign.jsonl")
+        dr = _count(BENCHMARK_ROOT / d / "direct.jsonl")
+        ind = _count(BENCHMARK_ROOT / d / "indirect.jsonl")
+        total = b + dr + ind
+        grand_total += total
+        table.add_row(d, str(b), str(dr), str(ind), str(total))
+
+    console.print(table)
+    console.print(f"[dim]{len(domains)} domains, {grand_total} tasks total[/dim]")
+
+
+@info_app.command("env", help="List Docker environments defined in dt_arena/config/env.yaml.")
+def info_env() -> None:
+    from utils.config import ENV_CONFIG_PATH
+
+    console = Console()
+    with ENV_CONFIG_PATH.open() as f:
+        cfg = yaml.safe_load(f) or {}
+    envs = cfg.get("environments", {}) or {}
+
+    table = Table(title="Docker environments", show_lines=False)
+    table.add_column("Name", style="bold cyan")
+    table.add_column("Ports", justify="right")
+    table.add_column("Max instances", justify="right")
+    table.add_column("Compose file")
+
+    for name in sorted(envs.keys()):
+        spec = envs[name] or {}
+        ports = len(spec.get("ports", {}) or {})
+        max_inst = spec.get("max_instances", "—")
+        compose = spec.get("docker_compose", "—")
+        table.add_row(name, str(ports), str(max_inst), str(compose))
+
+    console.print(table)
+    console.print(f"[dim]{len(envs)} environments[/dim]")
+
+
+# ── `dtap agent …` subcommands ───────────────────────────────────────────────
+
+_AGENT_BACKENDS = [
+    ("openaisdk",  "openai",    "OpenAI Agents SDK"),
+    ("claudesdk",  "claude",    "Anthropic Claude Agent SDK"),
+    ("googleadk",  "google",    "Google Agent Development Kit"),
+    ("langchain",  "langchain", "LangChain ReAct-style agent"),
+    ("pocketflow", "pocketflow", "Pocketflow ReAct agent"),
+    ("openclaw",   "claude",    "Claude SDK adapter (openclaw)"),
+]
+
+
+@agent_app.command("list", help="List supported agent backends and their pip extras.")
+def agent_list() -> None:
+    console = Console()
+    table = Table(title="Agent backends", show_lines=False, expand=False)
+    table.add_column("Agent type", style="bold cyan", no_wrap=True)
+    table.add_column("Install extra", no_wrap=True)
+    table.add_column("Description")
+
+    for agent_type, extra, desc in _AGENT_BACKENDS:
+        table.add_row(agent_type, escape(f"decodingtrust-agent-sdk[{extra}]"), desc)
+
+    console.print(table)
+    console.print("[dim]Pass --agent-type <name> to 'dtap eval'.[/dim]")
+
+
+@agent_app.command(
+    "create",
+    help="Scaffold a new custom agent under agent/<name>/.",
+)
+def agent_create(
+    name: str = typer.Argument(..., help="Agent name (e.g. my_agent). Becomes the directory name."),
+    framework: str = typer.Option(
+        None,
+        "--framework",
+        "-f",
+        help=f"Inherit from an existing backend ({', '.join(SUPPORTED_FRAMEWORKS)}). Omit to scaffold from scratch.",
+    ),
+    out: Path = typer.Option(
+        Path("agent"),
+        "--out",
+        help="Parent directory for the new scaffold (default: ./agent).",
+    ),
+    force: bool = typer.Option(
+        False,
+        "--force",
+        help="Overwrite an existing scaffold directory.",
+    ),
+) -> None:
+    console = Console()
+    try:
+        target = scaffold(name=name, framework=framework, out_root=out, force=force)
+    except (ValueError, FileExistsError) as e:
+        console.print(f"[red]error:[/red] {e}")
+        raise typer.Exit(code=1)
+
+    console.print(f"[green]Created[/green] {target}")
+    for f in sorted(target.iterdir()):
+        console.print(f"  [dim]·[/dim] {f.relative_to(target.parent)}")
+    console.print(
+        f"\nNext: edit [bold]{target}/agent.py[/bold], then run "
+        f"[cyan]python -m {target.name}.example[/cyan] (from {out})."
+    )

cli/scaffold.py

@@ -0,0 +1,334 @@
+"""Templates for `dtap agent create`."""
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Optional
+
+
+# kind="wrapper": user writes build_native_agent() returning a framework-native
+#                 agent; we wrap it via build_agent(native_agent=...).
+# kind="subclass": user subclasses the framework's DTap Agent class directly.
+
+_FRAMEWORK_INFO = {
+    "openaisdk": {
+        "kind": "wrapper",
+        "native_import": "from agents import Agent as OpenAIAgent",
+        "native_factory": (
+            'return OpenAIAgent(\n'
+            '        name="{name}",\n'
+            '        instructions="You are a helpful agent. TODO: customize.",\n'
+            '        tools=[],\n'
+            '    )'
+        ),
+    },
+    "pocketflow": {
+        "kind": "wrapper",
+        "native_import": "from agent.pocketflow.src.native_agent import NativeMCPReactAgent",
+        "native_factory": (
+            "# TODO: configure your NativeMCPReactAgent\n"
+            "    return NativeMCPReactAgent(\n"
+            '        system_prompt="You are a helpful agent.",\n'
+            "    )"
+        ),
+    },
+    "claudesdk": {
+        "kind": "subclass",
+        "base_import": "from agent.claudesdk import ClaudeSDKAgent",
+        "base_class": "ClaudeSDKAgent",
+    },
+    "googleadk": {
+        "kind": "subclass",
+        "base_import": "from agent.googleadk import GoogleADKAgent",
+        "base_class": "GoogleADKAgent",
+    },
+    "langchain": {
+        "kind": "subclass",
+        "base_import": "from agent.langchain import LangChainAgent",
+        "base_class": "LangChainAgent",
+    },
+    "openclaw": {
+        "kind": "subclass",
+        "base_import": "from agent.openclaw import OpenClawAgent",
+        "base_class": "OpenClawAgent",
+    },
+}
+
+SUPPORTED_FRAMEWORKS = sorted(_FRAMEWORK_INFO.keys())
+
+
+def _to_class_name(name: str) -> str:
+    """Turn 'my_custom-agent' into 'MyCustomAgent' (avoid doubling the 'Agent' suffix)."""
+    parts = re.split(r"[_\- ]+", name.strip())
+    camel = "".join(p.capitalize() for p in parts if p)
+    return camel if camel.endswith("Agent") else camel + "Agent"
+
+
+def _validate_name(name: str) -> None:
+    if not re.fullmatch(r"[a-zA-Z][a-zA-Z0-9_\-]*", name):
+        raise ValueError(
+            f"Invalid agent name '{name}'. Use letters, digits, underscores, hyphens, "
+            "starting with a letter."
+        )
+
+
+# ── Templates ────────────────────────────────────────────────────────────────
+
+_TPL_INIT_WRAPPER = '''"""Custom agent: {class_name}
+
+Generated by `dtap agent create {name} --framework {framework}`.
+"""
+
+from .agent import build_native_agent
+
+__all__ = ["build_native_agent"]
+'''
+
+_TPL_INIT_SUBCLASS = '''"""Custom agent: {class_name}
+
+Generated by `dtap agent create {name}{framework_arg}`.
+"""
+
+from .agent import {class_name}
+
+__all__ = ["{class_name}"]
+'''
+
+_TPL_AGENT_WRAPPER = '''"""Native-agent factory for {class_name} (wrapped at runtime by DTap)."""
+
+{native_import}
+
+
+def build_native_agent():
+    """Return your native {framework} agent.
+
+    DTap will wrap this via `build_agent(native_agent=...)`, which adds
+    red-teaming MCP servers and trajectory tracking automatically.
+    """
+    {native_factory}
+'''
+
+_TPL_AGENT_SUBCLASS = '''"""Custom DTap agent {class_name} (extends {base_class})."""
+
+from typing import Any, Dict, Optional
+
+{base_import}
+
+
+class {class_name}({base_class}):
+    """TODO: describe what {name} does differently from {base_class}."""
+
+    async def initialize(self) -> None:
+        await super().initialize()
+        # TODO: extra setup after MCP servers are connected.
+
+    async def run(self, user_input: str, metadata: Optional[Dict[str, Any]] = None) -> Any:
+        # TODO: customize prompt assembly / tool routing here if needed.
+        return await super().run(user_input, metadata=metadata)
+
+    async def cleanup(self) -> None:
+        # TODO: release any extra resources your agent owns.
+        await super().cleanup()
+'''
+
+_TPL_AGENT_SCRATCH = '''"""Custom DTap agent {class_name}.
+
+Subclasses the abstract Agent base. Fill in the four abstract methods.
+"""
+
+from typing import Any, Dict, Optional
+
+from dt_arena.src.types.agent import Agent, AgentConfig, MCPServerConfig, RuntimeConfig
+
+
+class {class_name}(Agent):
+    """TODO: describe what {name} does."""
+
+    def __init__(
+        self,
+        agent_config: Optional[AgentConfig] = None,
+        runtime_config: Optional[RuntimeConfig] = None,
+    ) -> None:
+        super().__init__(agent_config, runtime_config)
+        # TODO: initialise per-instance state (model client, tracing, etc.)
+
+    async def initialize(self) -> None:
+        # TODO: connect MCP servers, set up tracing, etc.
+        raise NotImplementedError("Implement initialize().")
+
+    def _create_mcp_server(self, server_config: MCPServerConfig) -> Any:
+        # TODO: return an SDK-specific MCP server instance for `server_config`.
+        raise NotImplementedError("Implement _create_mcp_server().")
+
+    async def run(self, user_input: str, metadata: Optional[Dict[str, Any]] = None) -> Any:
+        # TODO: invoke the underlying LLM/agent and return an AgentResult.
+        raise NotImplementedError("Implement run().")
+
+    async def cleanup(self) -> None:
+        # TODO: close MCP server connections and release resources.
+        raise NotImplementedError("Implement cleanup().")
+'''
+
+_TPL_EXAMPLE_WRAPPER = '''"""Smoke test for the {class_name} scaffold (wrapper pattern)."""
+
+import asyncio
+
+from dt_arena.src.types.agent import AgentConfig, RuntimeConfig
+from utils.agent_helpers import build_agent
+
+from . import build_native_agent
+
+
+async def main() -> None:
+    native = build_native_agent()
+    agent = build_agent(
+        native_agent=native,
+        agent_cfg=AgentConfig(system_prompt=""),
+        runtime_cfg=RuntimeConfig(model="gpt-4o", max_turns=10),
+    )
+    async with agent:
+        result = await agent.run("Say hello.", metadata={{"task_id": "smoke"}})
+        print(result.final_output)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+'''
+
+_TPL_EXAMPLE_SUBCLASS = '''"""Smoke test for the {class_name} scaffold."""
+
+import asyncio
+
+from dt_arena.src.types.agent import AgentConfig, RuntimeConfig
+
+from .agent import {class_name}
+
+
+async def main() -> None:
+    agent = {class_name}(
+        agent_config=AgentConfig(system_prompt=""),
+        runtime_config=RuntimeConfig(model="gpt-4o", max_turns=10),
+    )
+    async with agent:
+        result = await agent.run("Say hello.", metadata={{"task_id": "smoke"}})
+        print(result.final_output)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+'''
+
+_TPL_README_WRAPPER = '''# {class_name}
+
+Generated by `dtap agent create {name} --framework {framework}`.
+
+## Layout
+
+- `agent.py` — your `build_native_agent()` factory. Edit this.
+- `example.py` — runs a single turn through `build_agent(native_agent=...)`.
+
+## Run
+
+```bash
+python -m agent.{name}.example
+```
+
+## Plug into `dtap eval`
+
+To make `dtap eval --agent-type {name}` work, register this module in the
+agent registry (see `utils/agent_helpers.py:_get_agent_registry`) or invoke
+`build_agent(native_agent=build_native_agent(), ...)` from your own driver.
+'''
+
+_TPL_README_SUBCLASS = '''# {class_name}
+
+Generated by `dtap agent create {name}{framework_arg}`.
+
+## Layout
+
+- `agent.py` — your `{class_name}` class. Edit the TODOs.
+- `example.py` — instantiates the agent and runs one turn.
+
+## Run
+
+```bash
+python -m agent.{name}.example
+```
+
+## Plug into `dtap eval`
+
+To make `dtap eval --agent-type {name}` work, register `{class_name}` in the
+agent registry (see `utils/agent_helpers.py:_get_agent_registry`).
+'''
+
+
+def scaffold(
+    name: str,
+    framework: Optional[str],
+    out_root: Path,
+    *,
+    force: bool = False,
+) -> Path:
+    """Create <out_root>/<name>/ with templates. Returns the directory path."""
+    _validate_name(name)
+    if framework and framework not in _FRAMEWORK_INFO:
+        raise ValueError(
+            f"Unknown framework '{framework}'. Supported: {', '.join(SUPPORTED_FRAMEWORKS)}."
+        )
+
+    class_name = _to_class_name(name)
+    target = out_root / name
+    if target.exists() and not force:
+        raise FileExistsError(f"{target} already exists. Pass --force to overwrite.")
+    target.mkdir(parents=True, exist_ok=True)
+
+    framework_arg = f" --framework {framework}" if framework else ""
+
+    if framework is None:
+        (target / "__init__.py").write_text(_TPL_INIT_SUBCLASS.format(
+            class_name=class_name, name=name, framework_arg=framework_arg,
+        ))
+        (target / "agent.py").write_text(_TPL_AGENT_SCRATCH.format(
+            class_name=class_name, name=name,
+        ))
+        (target / "example.py").write_text(_TPL_EXAMPLE_SUBCLASS.format(
+            class_name=class_name,
+        ))
+        (target / "README.md").write_text(_TPL_README_SUBCLASS.format(
+            class_name=class_name, name=name, framework_arg=framework_arg,
+        ))
+        return target
+
+    info = _FRAMEWORK_INFO[framework]
+    if info["kind"] == "wrapper":
+        (target / "__init__.py").write_text(_TPL_INIT_WRAPPER.format(
+            class_name=class_name, name=name, framework=framework,
+        ))
+        (target / "agent.py").write_text(_TPL_AGENT_WRAPPER.format(
+            class_name=class_name, framework=framework,
+            native_import=info["native_import"],
+            native_factory=info["native_factory"].format(name=name),
+        ))
+        (target / "example.py").write_text(_TPL_EXAMPLE_WRAPPER.format(
+            class_name=class_name,
+        ))
+        (target / "README.md").write_text(_TPL_README_WRAPPER.format(
+            class_name=class_name, name=name, framework=framework,
+        ))
+    else:  # subclass
+        (target / "__init__.py").write_text(_TPL_INIT_SUBCLASS.format(
+            class_name=class_name, name=name, framework_arg=framework_arg,
+        ))
+        (target / "agent.py").write_text(_TPL_AGENT_SUBCLASS.format(
+            class_name=class_name, name=name,
+            base_import=info["base_import"],
+            base_class=info["base_class"],
+        ))
+        (target / "example.py").write_text(_TPL_EXAMPLE_SUBCLASS.format(
+            class_name=class_name,
+        ))
+        (target / "README.md").write_text(_TPL_README_SUBCLASS.format(
+            class_name=class_name, name=name, framework_arg=framework_arg,
+        ))
+
+    return target