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

agent/openaisdk/src/agent.py

@@ -0,0 +1,491 @@
+import os
+from datetime import datetime
+from typing import Dict, Any, Optional, List, Union
+
+from openai import AsyncOpenAI
+from agents import Agent as OpenAIAgent, Runner, ModelSettings, set_trace_processors, function_tool
+from agents import OpenAIChatCompletionsModel
+from agents.tracing import trace
+
+from dt_arena.src.types.agent import (
+    Agent,
+    AgentConfig,
+    RuntimeConfig,
+    MCPServerConfig,
+    AgentResult,
+)
+
+from .mcp_wrapper import MCPServerSse, MCPServerStreamableHttp, MCPServerStdio
+from dt_arena.src.types.trajectory import Trajectory
+
+from .utils import OpenAITraceProcessor, OpenAIAgentTrajectory
+from utils.skill_helpers import (
+    create_injected_skills_directory,
+    cleanup_temp_directory,
+    parse_skill_metadata,
+    load_skill_full_content,
+    scan_available_skills,
+)
+
+
+# Default MCP client read timeout for all transports.
+MCP_CLIENT_TIMEOUT_SECONDS: float = float(os.getenv("MCP_CLIENT_TIMEOUT_SECONDS", "30"))
+
+
+class OpenAISDKAgent(Agent):
+    """General Agent Workflow based on OpenAI Agents SDK"""
+
+    def __init__(
+        self,
+        agent_config: AgentConfig,
+        runtime_config: Optional[RuntimeConfig] = None,
+    ):
+        """
+        Initialize OpenAI SDK Agent
+
+        Args:
+            agent_config: Agent configuration (system_prompt and mcp_servers)
+            runtime_config: Runtime configuration (model, temperature, max_turns, output_dir)
+        """
+        super().__init__(agent_config, runtime_config)
+
+        # Setup output directory
+        output_dir = self.runtime_config.output_dir or os.path.join(
+            os.getcwd(), "results"
+        )
+        self.output_dir = output_dir
+        self.traces_dir = os.path.join(self.output_dir, "traces")
+        self.trajectories_dir = self.output_dir
+        os.makedirs(self.traces_dir, exist_ok=True)
+        os.makedirs(self.trajectories_dir, exist_ok=True)
+
+        # Initialize trace file and processor
+        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.trace_file = os.path.join(
+            self.traces_dir, f"traces_{self.timestamp}.jsonl"
+        )
+        self.trace_processor = OpenAITraceProcessor(self.trace_file)
+
+        # Use only local trace processor (replaces default OpenAI backend processor)
+        set_trace_processors([self.trace_processor])
+
+        # Initialize trajectory converter
+        self.trajectory_converter = OpenAIAgentTrajectory(
+            self.trajectories_dir, self.timestamp
+        )
+
+        # Multi-turn conversation support: stores last RunResult for to_input_list() to build conversation history
+        self._last_result = None
+
+        # Active trace context for multi-turn
+        self._active_trace = None
+        self._active_trace_id = None
+        self._trace_metadata = None
+
+        # Current execution result and trajectory
+        self._current_trajectory: Optional[Trajectory] = None
+        self._turn_count: int = 0
+
+        # Tool descriptions cache (for debug mode)
+        self._tool_descriptions: Dict[str, str] = {}
+
+        # Skill injection temp directory
+        self._skill_temp_dir: Optional[str] = None
+
+    async def _setup_skills(self) -> None:
+        """Setup skill directories and apply any skill injections."""
+        skill_directories = self.config.skill_directories if self.config else []
+        skill_injection = self.runtime_config.skill_injection
+
+        has_create_mode = skill_injection and any(
+            any(inj.mode == "create" for inj in injs)
+            for injs in skill_injection.values()
+        )
+
+        if not skill_directories and not has_create_mode:
+            return
+
+        self._skill_temp_dir = create_injected_skills_directory(
+            source_skill_dirs=skill_directories,
+            skill_injection=skill_injection,
+            skill_subpath=".agents/skills",
+            base_dir=self.output_dir,
+        )
+
+    def _build_tools_for_skill(self) -> list:
+        """Build a load_skill function tool that returns full SKILL.md content."""
+        tools = []
+        if not self._skill_temp_dir:
+            return tools
+
+        skills_dir = os.path.join(self._skill_temp_dir, ".agents", "skills")
+
+        @function_tool
+        def load_skill(skill_name: str) -> str:
+            """Load a skill by name and return its full instructions.
+
+            Use this tool to retrieve the complete content of a skill file.
+            After loading, follow the skill's instructions to complete the task.
+
+            Args:
+                skill_name: The name of the skill to load.
+
+            Returns:
+                The full content of the skill's SKILL.md file, or an error message if not found.
+            """
+            available = scan_available_skills(skills_dir)
+            for skill in available:
+                if skill.get("name") == skill_name:
+                    return load_skill_full_content(skill["path"])
+            names = [s.get("name", "") for s in available]
+            return f"Skill '{skill_name}' not found. Available skills: {names}"
+
+        tools.append(load_skill)
+        return tools
+
+    def _build_instructions_for_skill(self) -> str:
+        """Build agent instructions, appending skill hints if skills are configured."""
+        instructions = self.config.system_prompt
+        if not self._skill_temp_dir:
+            return instructions
+
+        skills_dir = os.path.join(self._skill_temp_dir, ".agents", "skills")
+        skill_entries = []
+        if os.path.isdir(skills_dir):
+            for item in sorted(os.listdir(skills_dir)):
+                skill_path = os.path.join(skills_dir, item)
+                if not os.path.isdir(skill_path):
+                    continue
+                skill_file = os.path.join(skill_path, "SKILL.md")
+                if not os.path.exists(skill_file):
+                    continue
+                metadata = parse_skill_metadata(skill_file)
+                if metadata:
+                    skill_entries.append(f"# Skill: {item}\n\n{metadata}")
+                else:
+                    skill_entries.append(f"# Skill: {item}")
+
+        if skill_entries:
+            skills_list = "\n".join(skill_entries)
+            instructions += (
+                "\n\n## Instructions for Using Skills\n"
+                "You have access to the `load_skill` tool that retrieves full skill instructions, "
+                "whose metadata are defined in the system prompt. "
+                "Based on the task requirements, if you need to use a skill, "
+                "call `load_skill` with the skill name to get its full instructions, then follow them.\n\n"
+                f"## Available skills:\n{skills_list}"
+            )
+        return instructions
+
+    async def initialize(self) -> None:
+        """Initialize agent and connect to MCP servers"""
+        if not self.config:
+            raise ValueError("Agent config is required")
+
+        # Setup skills if configured
+        await self._setup_skills()
+
+        # Load and connect to MCP servers
+        await self.load_mcp_servers()
+
+        # Build tools and instructions (with skill support)
+        tools = self._build_tools_for_skill()
+        instructions = self._build_instructions_for_skill()
+
+        # Determine model: use Chat Completions API for OpenRouter,
+        # since it only supports the Chat Completions format, not the
+        # Responses API that the SDK uses by default.
+        base_url = os.environ.get("OPENAI_BASE_URL", "")
+        if base_url and "openrouter.ai" in base_url:
+            model = OpenAIChatCompletionsModel(
+                model=self.runtime_config.model,
+                openai_client=AsyncOpenAI(
+                    base_url=base_url,
+                    api_key=os.environ.get("OPENAI_API_KEY"),
+                ),
+            )
+        else:
+            model = self.runtime_config.model
+
+        # Create OpenAI agent instance using runtime parameters
+        self.agent = OpenAIAgent(
+            name=self.config.name,
+            instructions=instructions,
+            model=model,
+            model_settings=ModelSettings(temperature=self.runtime_config.temperature)
+            if self.runtime_config.temperature is not None
+            else ModelSettings(),
+            mcp_servers=self.mcp_servers,
+            tools=tools,
+        )
+
+    def _create_mcp_server(self, server_config: MCPServerConfig) -> Any:
+        """
+        Create OpenAI SDK-specific MCP server instance.
+
+        If mcp_injection is configured in runtime_config, the wrapper classes
+        will automatically apply tool description injections.
+
+        Args:
+            server_config: Configuration for the MCP server
+
+        Returns:
+            OpenAI SDK MCP server instance (MCPServerSse, MCPServerStdio, or MCPServerStreamableHttp)
+        """
+        # Get tool injections for this server (if any)
+        tool_injections = None
+        if self.runtime_config.mcp_injection:
+            tool_injections = self.runtime_config.mcp_injection.get(server_config.name)
+
+        if server_config.transport == "sse":
+            return MCPServerSse(
+                name=server_config.name,
+                params={"url": server_config.url},
+                cache_tools_list=server_config.cache_tools_list,
+                client_session_timeout_seconds=MCP_CLIENT_TIMEOUT_SECONDS,
+                tool_injections=tool_injections,
+                hook_manager=self.hook_manager,
+            )
+        elif server_config.transport == "stdio":
+            params = {"command": server_config.command}
+            if server_config.args:
+                params["args"] = server_config.args
+            if server_config.env:
+                params["env"] = server_config.env
+            return MCPServerStdio(
+                name=server_config.name,
+                params=params,
+                cache_tools_list=server_config.cache_tools_list,
+                client_session_timeout_seconds=MCP_CLIENT_TIMEOUT_SECONDS,
+                tool_injections=tool_injections,
+                hook_manager=self.hook_manager,
+            )
+        elif server_config.transport == "http":
+            return MCPServerStreamableHttp(
+                name=server_config.name,
+                params={"url": server_config.url},
+                cache_tools_list=server_config.cache_tools_list,
+                client_session_timeout_seconds=MCP_CLIENT_TIMEOUT_SECONDS,
+                tool_injections=tool_injections,
+                hook_manager=self.hook_manager,
+            )
+        else:
+            raise ValueError(f"Unsupported transport type: {server_config.transport}")
+
+    async def load_mcp_servers(self) -> List[Any]:
+        """
+        Load and connect to MCP servers.
+
+        Override base class to add explicit connect() calls required by OpenAI SDK.
+        """
+        # Use base class to create servers and track names
+        await super().load_mcp_servers()
+
+        # OpenAI SDK requires explicit connection
+        for server in self.mcp_servers:
+            await server.connect()
+
+        # Cache tool descriptions for debug mode (with injections applied)
+        if self.runtime_config.debug:
+            await self._cache_tool_descriptions()
+
+        return self.mcp_servers
+
+    async def _cache_tool_descriptions(self) -> None:
+        """Cache tool descriptions from all MCP servers (for debug mode)."""
+        for i, server in enumerate(self.mcp_servers):
+            server_name = (
+                self.mcp_server_names[i]
+                if i < len(self.mcp_server_names)
+                else "unknown"
+            )
+            try:
+                # list_tools() returns tools with injections already applied
+                tools = await server.list_tools()
+                for tool in tools:
+                    key = f"{server_name}:{tool.name}"
+                    self._tool_descriptions[key] = tool.description or ""
+            except Exception as e:
+                print(
+                    f"[DEBUG] Failed to cache tool descriptions from {server_name}: {e}"
+                )
+
+    async def run(
+        self,
+        user_input: Union[str, List[str]],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> AgentResult:
+        """
+        Run the agent with given input. Supports multi-turn conversations.
+
+        Args:
+            user_input: User instruction/query. Can be:
+                - str: Single query (backward compatible)
+                - List[str]: Multiple queries processed sequentially with context preserved
+
+        metadata: Optional metadata (task_id, domain, category, instruction)
+
+        Returns:
+            AgentResult with final output, turn count, trace ID, and trajectory
+        """
+        if not self.agent:
+            raise RuntimeError(
+                "Agent not initialized. Call initialize() first or use async context manager."
+            )
+
+        # Normalize input to list for uniform processing
+        if isinstance(user_input, str):
+            user_inputs = [user_input]
+        else:
+            user_inputs = user_input
+
+        if self._trace_metadata is None:
+            self._trace_metadata = metadata or {}
+
+        # Set/update instruction in metadata
+        if "instruction" not in self._trace_metadata:
+            if len(user_inputs) == 1:
+                instr = user_inputs[0] or ""
+            else:
+                instr = " | ".join(
+                    user_inputs
+                )  # For multi-turn list input, combine queries with separator
+            self._trace_metadata["instruction"] = instr
+
+        # Include debug info in trace metadata if debug mode is enabled
+        if self.runtime_config.debug:
+            if self._tool_descriptions:
+                self._trace_metadata["tool_descriptions"] = self._tool_descriptions
+            if self.agent and hasattr(self.agent, "instructions"):
+                self._trace_metadata["system_prompt"] = self.agent.instructions
+
+        # Start new trace if none active
+        if self._active_trace is None:
+            self._active_trace = trace(
+                workflow_name=self._trace_metadata.get("task_id", "Agent Task"),
+                metadata=self._trace_metadata,
+            )
+            self._active_trace.__enter__()
+            self._active_trace_id = self._active_trace.trace_id
+
+        result = None
+
+        for i, query in enumerate(user_inputs):
+            # Build input for this turn
+            if self._last_result is not None:
+                # Use to_input_list() to get conversation history, then append new user message
+                input_items = self._last_result.to_input_list() + [
+                    {"role": "user", "content": query}
+                ]
+            else:
+                input_items = query  # first turn
+
+            # Run agent for this turn
+            result = await Runner.run(
+                starting_agent=self.agent,
+                input=input_items,
+                max_turns=self.runtime_config.max_turns,
+            )
+
+            # Store result for next turn's history
+            self._last_result = result
+
+            # Record response after each turn (including intermediate responses)
+            if result:
+                self.trace_processor.write_agent_response(
+                    self._active_trace_id, str(result.final_output)
+                )
+
+            self._turn_count += 1
+
+        # Close trace context before generating trajectory
+        if self._active_trace is not None:
+            try:
+                self._active_trace.__exit__(None, None, None)
+            except Exception as e:
+                print(f"[WARNING] Failed to close trace context: {e}")
+            self._active_trace = None
+
+        # Generate trajectory at end of run
+        self._generate_trajectory(result.final_output if result else None)
+
+        return self.get_result()
+
+    def _generate_trajectory(self, final_output: Optional[str] = None) -> None:
+        """Generate trajectory from current trace file"""
+        try:
+            if os.path.exists(self.trace_file):
+                trajectories = self.trajectory_converter.process_trace_file(
+                    self.trace_file, output_name="trajectory"
+                )
+                if trajectories:
+                    self._current_trajectory = trajectories[
+                        -1
+                    ]  # Get the latest trajectory
+        except Exception as e:
+            print(f"[WARNING] Failed to generate trajectory: {e}")
+
+    def get_result(self) -> Optional[AgentResult]:
+        """
+        Get the current execution result.
+
+        Returns:
+            AgentResult with final output, turn count, trace ID, and trajectory.
+            Returns None if no execution has been performed yet.
+        """
+        if self._last_result is None:
+            return None
+
+        return AgentResult(
+            final_output=str(self._last_result.final_output)
+            if self._last_result.final_output
+            else None,
+            turn_count=self._turn_count,
+            trajectory=self._current_trajectory,
+            trace_id=self._active_trace_id or "",
+        )
+
+    def reset_conversation(self) -> None:
+        """
+        Reset conversation history for a new session.
+
+        Call this to clear the accumulated context and start fresh.
+        Automatically called during cleanup().
+
+        This will:
+        - Close the active trace context (finalizing the trajectory)
+        - Clear conversation history
+        - Allow a new trace to be started on next run() call
+        """
+        # Close active trace context if exists
+        if self._active_trace is not None:
+            try:
+                self._active_trace.__exit__(None, None, None)
+            except Exception as e:
+                print(f"[WARNING] Failed to close trace context: {e}")
+            self._active_trace = None
+            self._active_trace_id = None
+
+        # Clear conversation state
+        self._last_result = None
+        self._trace_metadata = None
+        self._current_trajectory = None
+        self._turn_count = 0
+
+    async def cleanup(self) -> None:
+        """Clean up resources and process traces"""
+        # Reset conversation history
+        self.reset_conversation()
+
+        # Cleanup MCP server connections
+        # Note: Cancel scope errors from anyio are suppressed in the MCP wrapper classes
+        for server in self.mcp_servers:
+            try:
+                await server.cleanup()
+            except BaseException as e:
+                print(f"[WARNING] Failed to cleanup server: {e}")
+
+        # Clean up skill temp directory
+        if self._skill_temp_dir:
+            cleanup_temp_directory(self._skill_temp_dir)
+            self._skill_temp_dir = None

agent/openaisdk/src/agent_wrapper.py

@@ -0,0 +1,143 @@
+import os
+from typing import Any, List, Optional
+
+from agents import Agent as OpenAIAgent
+
+from dt_arena.src.types.agent import AgentConfig, RuntimeConfig
+
+from .agent import OpenAISDKAgent
+
+
+class OpenAISDKNativeWrapper(OpenAISDKAgent):
+    """
+    Wraps a pre-built OpenAI SDK Agent for evaluation.
+
+
+    Overrides:
+    - __init__: Accept native agent instead of building from config
+    - initialize(): Clone native agent and add red-teaming MCP servers
+    """
+
+    def __init__(
+        self,
+        native_agent: OpenAIAgent,
+        agent_config: Optional[AgentConfig] = None,
+        runtime_config: Optional[RuntimeConfig] = None,
+    ):
+        """
+        Initialize the wrapper around a native OpenAI SDK Agent.
+
+        Args:
+            native_agent: Pre-built agents.Agent instance from user
+            agent_config: AgentConfig with red-teaming MCP servers to add
+            runtime_config: RuntimeConfig with model settings, output_dir, mcp_injection
+        """
+        # Store the native agent before calling parent init
+        self._native_agent = native_agent
+
+        # Track red-teaming MCP servers we add (for cleanup)
+        self._redteaming_mcp_servers: List[Any] = []
+
+        # Set up tracing, directories, trajectory converter
+        super().__init__(
+            agent_config=agent_config or AgentConfig(system_prompt=""),
+            runtime_config=runtime_config,
+        )
+
+    async def initialize(self) -> None:
+        """
+        Initialize the wrapper.
+
+        The native agent's existing MCP servers are assumed to be already
+        connected by the user.
+        """
+        self.agent = self._native_agent.clone()
+
+        # Create a new mcp_servers list to avoid modifying the original
+        original_servers = list(self.agent.mcp_servers or [])
+        self.agent.mcp_servers = original_servers
+
+        # Add red-teaming MCP servers from config (i.e., with tool injection)
+        await self._add_redteaming_mcp_servers()
+
+        # Setup skills if configured
+        await self._setup_skills()
+        if self._skill_temp_dir:
+            tools = self._build_tools()
+            if tools:
+                existing_tools = list(self.agent.tools or [])
+                existing_tools.extend(tools)
+                self.agent.tools = existing_tools
+
+    async def _add_redteaming_mcp_servers(self) -> None:
+        """
+        Add red-teaming MCP servers from agent_config to the agent.
+
+        Tool injections from runtime_config.mcp_injection are applied to these
+        servers via the parent class's _create_mcp_server() method.
+        """
+        if not self.config or not self.config.mcp_servers:
+            return
+
+        for server_config in self.config.mcp_servers:
+            if not server_config.enabled:
+                continue
+
+            # Check if server with same name already exists in native agent
+            existing_names = {getattr(s, 'name', None) for s in (self.agent.mcp_servers or [])}
+            if server_config.name in existing_names:
+                print(f"[WARNING] MCP server '{server_config.name}' already exists in native agent. "
+                      f"Red-teaming server will be added with potential name conflict.")
+
+            # Create red-teaming server using parent's method (applies tool injection)
+            server = self._create_mcp_server(server_config)
+
+            # Connect the server
+            await server.connect()
+
+            # Track for cleanup
+            self._redteaming_mcp_servers.append(server)
+            self.mcp_servers.append(server)
+            self.mcp_server_names.append(server_config.name)
+
+            # Add to agent's mcp_servers list
+            self.agent.mcp_servers.append(server)
+
+    async def cleanup(self) -> None:
+        """
+        Clean up resources.
+
+        Cleans up all MCP servers (both native agent's and red-teaming servers).
+        """
+        # Reset conversation state
+        self.reset_conversation()
+
+        # Clean up skill temp directory
+        if self._skill_temp_dir:
+            from utils.skill_helpers import cleanup_temp_directory
+            cleanup_temp_directory(self._skill_temp_dir)
+            self._skill_temp_dir = None
+
+        # Cleanup all MCP servers on the cloned agent
+        # This includes both the native agent's servers and red-teaming servers we added
+        if self.agent is not None:
+            for server in (self.agent.mcp_servers or []):
+                try:
+                    await server.cleanup()
+                except BaseException as e:
+                    print(f"[WARNING] Failed to cleanup server: {e}")
+            self.agent.mcp_servers = []
+
+        self._redteaming_mcp_servers = []
+        self.mcp_servers = []
+        self.mcp_server_names = []
+
+    @property
+    def native_agent(self) -> OpenAIAgent:
+        """Access the original (unmodified) native agent."""
+        return self._native_agent
+
+    @property
+    def wrapped_agent(self) -> OpenAIAgent:
+        """Access the cloned and modified agent used for execution."""
+        return self.agent