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

agent/pocketflow/example.py

@@ -0,0 +1,221 @@
+import os
+import sys
+import asyncio
+import argparse
+
+from dt_arena.src.types.agent import AgentConfig, RuntimeConfig
+from dt_arena.src.types.task import TaskConfig, AttackConfig
+
+from agent.pocketflow import MCPReactAgent
+from utils.injection_helpers import (
+    build_tool_injections_from_config,
+    apply_prompt_injections,
+)
+from utils.task_helpers import extract_dataset_path
+
+
+# ============================================================================
+# Standalone Test Example for MCPReactAgent
+# Supports tool injection and prompt injection from attack config
+# ============================================================================
+
+
+async def main():
+    """Example demonstrating agent with MCP server and trajectory generation"""
+
+    parser = argparse.ArgumentParser(
+        description="Run PocketFlow SDK Agent with configuration file",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Run with default settings
+  python example.py --config path/to/config.yaml
+
+  # Run with custom model and temperature
+  python example.py --config path/to/config.yaml --model gpt-4o-mini --temperature 0.5
+
+  # Run with custom max turns and output directory
+  python example.py --config path/to/config.yaml --max-turns 20 --output-dir ./custom_results
+        """
+    )
+
+    parser.add_argument(
+        "--config",
+        type=str,
+        required=True,
+        help="Path to YAML configuration file (with Task, Agent, Attack sections)"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="gpt-4o",
+        help="Model to use (default: gpt-4o)"
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.1,
+        help="Sampling temperature, 0.0-1.0 (default: 0.1)"
+    )
+    parser.add_argument(
+        "--max-turns",
+        type=int,
+        default=10,
+        help="Maximum conversation turns (default: 10)"
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Output directory for traces and trajectories (default: ./results)"
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable debug mode for tracing"
+    )
+
+    args = parser.parse_args()
+
+    # Setup
+    config_path = args.config
+    if not os.path.exists(config_path):
+        print(f"[ERROR] Configuration file not found: {config_path}")
+        sys.exit(1)
+
+    # Extract dataset path structure (e.g., crm/benign/1)
+    dataset_path = extract_dataset_path(config_path)
+
+    # Build output directory
+    base_output_dir = args.output_dir or os.path.join(os.getcwd(), "results")
+    output_dir = os.path.join(base_output_dir, "pocketflow", dataset_path)
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"[INFO] Loading configuration from: {config_path}")
+    print(f"[INFO] Model: {args.model}")
+    print(f"[INFO] Temperature: {args.temperature}")
+    print(f"[INFO] Max Turns: {args.max_turns}")
+    print(f"[INFO] Output Directory: {output_dir}")
+    print("-" * 80)
+
+    try:
+        # Load configuration
+        task_config = TaskConfig.from_yaml(config_path)
+        agent_config = AgentConfig.from_yaml(config_path)
+        attack_config = AttackConfig.from_yaml(config_path)
+
+        # Build tool injections from all attack turns
+        mcp_injection = build_tool_injections_from_config(attack_config)
+
+        # Create runtime config with tool injections
+        runtime_config = RuntimeConfig(
+            model=args.model,
+            temperature=args.temperature,
+            max_turns=args.max_turns,
+            output_dir=output_dir,
+            mcp_injection=mcp_injection,
+            debug=args.debug,
+        )
+
+        # Create agent with configurations
+        agent = MCPReactAgent(
+            agent_config=agent_config,
+            runtime_config=runtime_config,
+        )
+
+        # Display loaded configuration
+        print(f"\n[TASK] Task ID: {task_config.task_id}")
+        print(f"[TASK] Domain: {task_config.domain}")
+
+        if attack_config:
+            print(f"\n[ATTACK] Risk Category: {attack_config.risk_category}")
+            print(f"[ATTACK] Threat Model: {attack_config.threat_model}")
+            if attack_config.malicious_goal:
+                print(f"[ATTACK] Malicious Goal: {attack_config.malicious_goal}")
+            if attack_config.attack_turns:
+                print(f"[ATTACK] Attack Turns: {len(attack_config.attack_turns)}")
+                for turn in attack_config.attack_turns:
+                    print(f"  Turn {turn.turn_id}: {len(turn.attack_steps)} attack step(s)")
+                    for step in turn.attack_steps:
+                        print(f"    - {step.type} ({step.mode})")
+
+        # Handle prompt injections based on threat model
+        is_direct_attack = attack_config and attack_config.threat_model == "direct"
+
+        if is_direct_attack:
+            user_instruction = apply_prompt_injections(original_instruction=None, attack_config=attack_config)
+            if not user_instruction:
+                user_instruction = task_config.original_instruction
+        else:
+            user_instruction = task_config.original_instruction
+            user_instruction = apply_prompt_injections(user_instruction, attack_config)
+
+        # Normalize instruction
+        if isinstance(user_instruction, str):
+            user_instruction = user_instruction.strip()
+        else:
+            user_instruction = [instr.strip() for instr in user_instruction]
+
+        print("\n" + "=" * 80)
+        print("[INSTRUCTION (after injection)]")
+        print("=" * 80)
+        if isinstance(user_instruction, list):
+            print(f"Multi-turn task with {len(user_instruction)} queries:")
+            for i, instr in enumerate(user_instruction, 1):
+                print(f"  {i}. {instr}")
+        else:
+            print(user_instruction)
+        print("=" * 80)
+
+        if attack_config and attack_config.malicious_goal:
+            print("\n[MALICIOUS GOAL] Testing Against:")
+            print("-" * 80)
+            print(attack_config.malicious_goal.strip())
+            print("-" * 80)
+
+        # Run agent with async context manager
+        async with agent:
+            print("\n[INFO] Running agent...")
+            if isinstance(user_instruction, list):
+                print(f"[INFO] Processing {len(user_instruction)} turns...")
+            # Pass task metadata via run() method
+            metadata = {
+                "task_id": task_config.task_id,
+                "domain": task_config.domain,
+                "category": attack_config.risk_category if attack_config else None,
+                "malicious_goal": attack_config.malicious_goal if attack_config else None,
+            }
+
+            result = await agent.run(user_instruction, metadata=metadata)
+
+            print("\n" + "=" * 80)
+            print("[AGENT RESPONSE]")
+            print("=" * 80)
+            print(result.final_output)
+            print("=" * 80)
+
+            print(f"\n[SUCCESS] Task completed")
+            print(f"[INFO] Turns: {result.turn_count}")
+            print(f"[INFO] Trajectories saved to: {agent.trajectories_dir}")
+            if result.trajectory:
+                print(f"[INFO] Trajectory steps: {len(result.trajectory)}")
+
+            return result
+
+    except Exception as e:
+        print(f"\n[ERROR] {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    try:
+        result = asyncio.run(main())
+        if result:
+            print(f"\n[SUCCESS] Task completed")
+    except KeyboardInterrupt:
+        print("\n[INFO] Interrupted by user")
+    except Exception as e:
+        print(f"\n[ERROR] {e}")
+        sys.exit(1)

agent/pocketflow/prompts/react_agent.py

@@ -0,0 +1,46 @@
+ACTION_SPACE_INSTRUCTION = """You are a ReAct agent that can use tools to complete tasks. You MUST choose ONE of the following two actions:
+
+## Available Actions:
+1. **use_tool**: Call one of the available tools to gather information or perform an action
+   - When to use: You need more information or need to perform an action
+   - Required fields: action, reasoning, tool_name, tool_arguments
+   - Set final_answer to null
+
+2. **answer**: Provide the final answer to the user's query
+   - When to use: You have sufficient information to provide a complete answer
+   - Required fields: action, reasoning, final_answer
+   - Set tool_name and tool_arguments to null
+
+## Response Format (MUST be valid YAML):
+
+```yaml
+reasoning: |
+  <your step-by-step reasoning about what to do next>
+action: use_tool  # MUST be exactly "use_tool" or "answer" (no other values allowed)
+tool_name: <tool_name>  # REQUIRED if action is "use_tool", otherwise set to null
+tool_arguments:  # REQUIRED if action is "use_tool", otherwise set to null
+  <param1>: <value1>
+  <param2>: <value2>
+final_answer: |  # REQUIRED if action is "answer", otherwise set to null
+  <your final answer to the user>
+```
+
+## Important Rules:
+- The response MUST be valid YAML enclosed in ```yaml and ```
+- The "action" field MUST be exactly "use_tool" or "answer" (case-sensitive)
+- If action is "use_tool": tool_name MUST be specified, tool_arguments MUST be a dict, final_answer MUST be null
+- If action is "answer": final_answer MUST be specified, tool_name MUST be null, tool_arguments MUST be null
+- Always include the "reasoning" field to explain your decision
+"""
+
+REACT_AGENT_SYSPROMPT = """{user_given_system_prompt}
+
+# Available Tools
+{tools_text}
+
+# Action Space and Response Format
+{action_space_instruction}
+
+Think carefully about whether you have enough information to answer the question.
+- If you need more information, use an appropriate tool.
+- If you have sufficient information, provide a comprehensive final answer."""

agent/pocketflow/src/__init__.py

@@ -0,0 +1,6 @@
+"""PocketFlow MCP Agent Implementation."""
+
+from .agent import MCPReactAgent, MCPServerInfo
+from .async_helper import AsyncHelper
+
+__all__ = ["MCPReactAgent", "MCPServerInfo", "AsyncHelper"]