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

agent/langchain/src/utils.py

@@ -0,0 +1,433 @@
+import os
+import json
+import uuid
+import threading
+from datetime import datetime, timezone
+from typing import List, Dict, Any, Optional
+from collections import defaultdict
+
+# Add parent directory to path to import types
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../..'))
+from dt_arena.src.types.trajectory import Trajectory
+
+
+class LangChainTraceProcessor:
+    """
+    Trace processor for LangChain agent that captures all execution data.
+    """
+    
+    def __init__(self, path: str):
+        self.path = path
+        self._lock = threading.Lock()
+        self._traces: Dict[str, Dict[str, Any]] = {}
+    
+    def _write(self, obj: Dict[str, Any]):
+        """Write a record to the trace file"""
+        with self._lock:
+            try:
+                with open(self.path, "a", encoding="utf-8") as f:
+                    f.write(json.dumps(obj, ensure_ascii=False, default=str) + "\n")
+            except Exception as e:
+                print(f"[WARNING] Failed to write trace: {e}")
+    
+    def start_trace(self, metadata: Dict[str, Any]) -> str:
+        """
+        Start a new trace
+        
+        Args:
+            metadata: Trace metadata (task_id, domain, category, etc.)
+            
+        Returns:
+            trace_id: Unique identifier for this trace
+        """
+        trace_id = str(uuid.uuid4())
+        
+        self._traces[trace_id] = {
+            "trace_id": trace_id,
+            "metadata": metadata,
+            "started_at": datetime.now(timezone.utc).isoformat(),
+            "steps": [],
+            "user_inputs": [],
+            "tool_calls": [],
+            "errors": []
+        }
+        
+        self._write({
+            "record": "trace_start",
+            "trace_id": trace_id,
+            "metadata": metadata,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+        
+        return trace_id
+    
+    def record_user_input(self, trace_id: str, user_input: str):
+        """Record user input message"""
+        if trace_id in self._traces:
+            self._traces[trace_id]["user_inputs"].append(user_input)
+        
+        self._write({
+            "record": "user_input",
+            "trace_id": trace_id,
+            "content": user_input,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+    
+    def record_tool_call(
+        self,
+        trace_id: str,
+        tool_name: str,
+        tool_input: Dict[str, Any],
+        tool_output: str,
+        server: Optional[str] = None
+    ):
+        """Record a tool call and its result"""
+        tool_call_data = {
+            "tool_name": tool_name,
+            "tool_input": tool_input,
+            "tool_output": tool_output,
+            "server": server,
+            "ts": datetime.now(timezone.utc).isoformat()
+        }
+        
+        if trace_id in self._traces:
+            self._traces[trace_id]["tool_calls"].append(tool_call_data)
+        
+        self._write({
+            "record": "tool_call",
+            "trace_id": trace_id,
+            **tool_call_data
+        })
+    
+    def record_agent_response(self, trace_id: str, response: str):
+        """Record agent's final response"""
+        self._write({
+            "record": "agent_response",
+            "trace_id": trace_id,
+            "response": response,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+
+    def write_agent_response(self, trace_id: str, response: str):
+        """
+        Write intermediate agent response (for multi-turn conversations).
+        
+        This is called after each turn to record the intermediate response,
+        distinct from record_agent_response which is for the final response.
+        """
+        self._write({
+            "record": "agent_turn_response",
+            "trace_id": trace_id,
+            "response": response,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+    
+    def record_error(self, trace_id: str, error: str):
+        """Record an error"""
+        if trace_id in self._traces:
+            self._traces[trace_id]["errors"].append(error)
+        
+        self._write({
+            "record": "error",
+            "trace_id": trace_id,
+            "error": error,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+    
+    def end_trace(self, trace_id: str, final_output: Optional[str], duration: float):
+        """End the trace and record final state"""
+        ended_at = datetime.now(timezone.utc).isoformat()
+        
+        if trace_id in self._traces:
+            self._traces[trace_id]["ended_at"] = ended_at
+            self._traces[trace_id]["duration"] = duration
+            self._traces[trace_id]["final_output"] = final_output
+        
+        self._write({
+            "record": "trace_end",
+            "trace_id": trace_id,
+            "final_output": final_output,
+            "duration": duration,
+            "ended_at": ended_at,
+            "ts": datetime.now(timezone.utc).isoformat()
+        })
+
+
+class LangChainTrajectoryConverter:
+    """Convert LangChain agent traces to standard trajectory format"""
+
+    def __init__(self, output_dir: str):
+        self.output_dir = output_dir
+        os.makedirs(self.output_dir, exist_ok=True)
+    
+    def load_traces(self, trace_file: str) -> Dict[str, Dict[str, Any]]:
+        """
+        Load traces from JSONL file
+        
+        Args:
+            trace_file: Path to trace JSONL file
+            
+        Returns:
+            Dictionary mapping trace_id to trace data
+        """
+        if not os.path.exists(trace_file):
+            raise FileNotFoundError(f"Trace file not found: {trace_file}")
+        
+        traces: Dict[str, Dict[str, Any]] = defaultdict(lambda: {
+            "metadata": {},
+            "user_inputs": [],
+            "tool_calls": [],
+            "agent_responses": [],
+            "turn_responses": [],  # Intermediate responses for multi-turn
+            "errors": [],
+            "final_output": None,
+            "duration": 0.0
+        })
+        
+        with open(trace_file, 'r', encoding='utf-8') as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                    
+                try:
+                    data = json.loads(line)
+                    record_type = data.get("record")
+                    trace_id = data.get("trace_id")
+                    
+                    if not trace_id:
+                        continue
+                    
+                    if record_type == "trace_start":
+                        traces[trace_id]["metadata"] = data.get("metadata", {})
+                        traces[trace_id]["started_at"] = data.get("ts")
+                        
+                    elif record_type == "user_input":
+                        traces[trace_id]["user_inputs"].append(data.get("content", ""))
+                        
+                    elif record_type == "tool_call":
+                        traces[trace_id]["tool_calls"].append({
+                            "tool_name": data.get("tool_name"),
+                            "tool_input": data.get("tool_input", {}),
+                            "tool_output": data.get("tool_output", ""),
+                            "server": data.get("server"),
+                            "ts": data.get("ts")
+                        })
+                        
+                    elif record_type == "agent_response":
+                        traces[trace_id]["agent_responses"].append(data.get("response", ""))
+                    
+                    elif record_type == "agent_turn_response":
+                        # Intermediate responses from multi-turn conversations
+                        traces[trace_id]["turn_responses"].append({
+                            "response": data.get("response", ""),
+                            "ts": data.get("ts")
+                        })
+                        
+                    elif record_type == "error":
+                        traces[trace_id]["errors"].append(data.get("error", ""))
+                        
+                    elif record_type == "trace_end":
+                        traces[trace_id]["final_output"] = data.get("final_output")
+                        traces[trace_id]["duration"] = data.get("duration", 0.0)
+                        traces[trace_id]["ended_at"] = data.get("ended_at")
+                        
+                except json.JSONDecodeError as e:
+                    print(f"[WARNING] Failed to parse trace line: {e}")
+                    continue
+        
+        return dict(traces)
+    
+    def convert_trace_to_trajectory(
+        self,
+        trace_id: str,
+        trace_data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Convert a single trace to trajectory format
+        
+        Args:
+            trace_id: Trace identifier
+            trace_data: Trace data dictionary
+            
+        Returns:
+            Trajectory dictionary
+        """
+        metadata = trace_data.get("metadata", {})
+        
+        # Create Trajectory instance
+        traj = Trajectory(
+            task_id=metadata.get("task_id", trace_id[:16]),
+            original_instruction=metadata.get("instruction", ""),
+            malicious_instruction=metadata.get("malicious_goal"),
+            domain=metadata.get("domain"),
+            risk_category=metadata.get("category")
+        )
+        
+        # Set duration and timestamp
+        traj.data["traj_info"]["duration"] = trace_data.get("duration", 0.0)
+        if trace_data.get("ended_at"):
+            traj.data["traj_info"]["timestamp"] = trace_data["ended_at"]
+        
+        # Build timeline of events for proper interleaving
+        # For multi-turn conversations, we need to interleave:
+        # user_input -> tool_calls -> agent_response -> user_input -> ...
+        
+        user_inputs = trace_data.get("user_inputs", [])
+        tool_calls = trace_data.get("tool_calls", [])
+        turn_responses = trace_data.get("turn_responses", [])
+        
+        # Track seen responses to avoid duplicates
+        seen_responses = set()
+        
+        # Process in order: for each user input, add related tool calls and response
+        # Note: This is a simplified approach assuming tool calls happen between user inputs
+        tool_call_idx = 0
+        turn_response_idx = 0
+        
+        for i, user_input in enumerate(user_inputs):
+            # Add user input
+            traj.append_user_step(user_input)
+            
+            # Add tool calls that happened after this user input
+            # (before next user input or end)
+            while tool_call_idx < len(tool_calls):
+                tool_call = tool_calls[tool_call_idx]
+                tool_name = tool_call.get("tool_name", "unknown")
+                tool_input = tool_call.get("tool_input", {})
+                tool_output = tool_call.get("tool_output", "")
+                server = tool_call.get("server")
+                
+                # Build action string
+                action_str = f"{tool_name}("
+                params = []
+                for k, v in tool_input.items():
+                    if isinstance(v, str):
+                        params.append(f'{k}="{v}"')
+                    else:
+                        params.append(f'{k}={v}')
+                action_str += ", ".join(params) + ")"
+                
+                # Add agent action step
+                traj.append_agent_step(
+                    action=action_str,
+                    tool_name=tool_name,
+                    tool_params=tool_input,
+                    server=server
+                )
+                
+                # Add tool return step
+                try:
+                    parsed_output = json.loads(tool_output) if isinstance(tool_output, str) else tool_output
+                except json.JSONDecodeError:
+                    parsed_output = tool_output
+                
+                traj.append_tool_return(
+                    result=parsed_output,
+                    tool_name=tool_name,
+                    server=server
+                )
+                
+                tool_call_idx += 1
+                
+                # Check if there's a turn response after this tool call
+                # Simple heuristic: if we've processed enough tool calls for this turn
+                if turn_response_idx < len(turn_responses):
+                    # Check if next item is a user input or more tool calls
+                    # For now, add turn response after tool calls
+                    break
+            
+            # Add turn response after tool calls for this turn
+            if turn_response_idx < len(turn_responses):
+                turn_resp = turn_responses[turn_response_idx]
+                response_text = turn_resp.get("response", "")
+                if response_text and response_text not in seen_responses:
+                    traj.append_agent_step(
+                        action="send_message_to_user",
+                        metadata={"message": response_text}
+                    )
+                    seen_responses.add(response_text)
+                turn_response_idx += 1
+        
+        # Process any remaining tool calls
+        while tool_call_idx < len(tool_calls):
+            tool_call = tool_calls[tool_call_idx]
+            tool_name = tool_call.get("tool_name", "unknown")
+            tool_input = tool_call.get("tool_input", {})
+            tool_output = tool_call.get("tool_output", "")
+            server = tool_call.get("server")
+            
+            action_str = f"{tool_name}("
+            params = []
+            for k, v in tool_input.items():
+                if isinstance(v, str):
+                    params.append(f'{k}="{v}"')
+                else:
+                    params.append(f'{k}={v}')
+            action_str += ", ".join(params) + ")"
+            
+            traj.append_agent_step(
+                action=action_str,
+                tool_name=tool_name,
+                tool_params=tool_input,
+                server=server
+            )
+            
+            try:
+                parsed_output = json.loads(tool_output) if isinstance(tool_output, str) else tool_output
+            except json.JSONDecodeError:
+                parsed_output = tool_output
+            
+            traj.append_tool_return(
+                result=parsed_output,
+                tool_name=tool_name,
+                server=server
+            )
+            
+            tool_call_idx += 1
+        
+        # Add final agent response (if different from turn responses)
+        final_output = trace_data.get("final_output")
+        if final_output and final_output not in seen_responses:
+            traj.append_agent_step(
+                action="send_message_to_user",
+                metadata={"message": final_output}
+            )
+        
+        return traj.to_dict()
+    
+    def process_trace_file(
+        self,
+        trace_file: str,
+        output_name: Optional[str] = None
+    ) -> List[str]:
+        """
+        Process trace file and generate trajectory files
+        
+        Args:
+            trace_file: Path to trace JSONL file
+            output_name: Optional prefix for output file names
+            
+        Returns:
+            List of saved trajectory file paths
+        """
+        traces = self.load_traces(trace_file)
+        saved_files = []
+        
+        for trace_id, trace_data in traces.items():
+            task_id = trace_data.get("metadata", {}).get("task_id")
+            
+            trajectory = self.convert_trace_to_trajectory(trace_id, trace_data)
+            
+            # Generate filename
+            filename = f"{output_name}_{task_id or trace_id[:8]}.json" if output_name else f"trajectory_{task_id or trace_id[:8]}.json"
+            output_path = os.path.join(self.output_dir, filename)
+            
+            # Save trajectory
+            with open(output_path, 'w', encoding='utf-8') as f:
+                json.dump(trajectory, f, indent=2, ensure_ascii=False)
+            
+            saved_files.append(output_path)
+        
+        return saved_files
+

agent/openaisdk/__init__.py

@@ -0,0 +1,17 @@
+"""
+OpenAI SDK Agent Implementation
+
+This module provides an agent implementation using the OpenAI Agents SDK
+with trajectory tracking and MCP server support.
+"""
+
+from .src.agent import OpenAISDKAgent
+from .src.agent_wrapper import OpenAISDKNativeWrapper
+from .src.utils import OpenAITraceProcessor, OpenAIAgentTrajectory
+
+__all__ = [
+    "OpenAISDKAgent",
+    "OpenAISDKNativeWrapper",
+    "OpenAITraceProcessor",
+    "OpenAIAgentTrajectory",
+]

agent/openaisdk/example.py

@@ -0,0 +1,228 @@
+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.openaisdk import OpenAISDKAgent
+from utils.injection_helpers import (
+    build_tool_injections_from_config,
+    build_skill_injections_from_config,
+    apply_prompt_injections,
+)
+from utils.task_helpers import extract_dataset_path
+
+
+# ============================================================================
+# Standalone Test Example for OpenAISDKAgent
+# 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 OpenAI 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, "openaisdk", 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)
+
+        # Build skill injections from all attack turns
+        skill_injection = build_skill_injections_from_config(attack_config)
+
+        # Create runtime config with tool and skill injections
+        runtime_config = RuntimeConfig(
+            model=args.model,
+            temperature=args.temperature,
+            max_turns=args.max_turns,
+            output_dir=output_dir,
+            mcp_injection=mcp_injection,
+            skill_injection=skill_injection,
+            debug=args.debug,
+        )
+
+        # Create agent with configurations
+        agent = OpenAISDKAgent(
+            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] Trace ID: {result.trace_id}")
+            print(f"[INFO] Traces saved to: {agent.traces_dir}")
+            print(f"[INFO] Trajectories saved to: {agent.trajectories_dir}")
+            if result.trajectory:
+                print(f"[INFO] Trajectory steps: {len(result.trajectory.data['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/openaisdk/src/__init__.py

@@ -0,0 +1,12 @@
+"""
+OpenAI SDK Agent source modules
+"""
+
+from .agent import OpenAISDKAgent
+from .utils import OpenAITraceProcessor, OpenAIAgentTrajectory
+
+__all__ = [
+    "OpenAISDKAgent",
+    "OpenAITraceProcessor",
+    "OpenAIAgentTrajectory",
+]