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

agent/googleadk/src/utils.py

@@ -0,0 +1,602 @@
+import os
+import json
+import threading
+from datetime import datetime, timezone
+from typing import List, Dict, Any, Optional
+from collections import defaultdict
+
+from dt_arena.src.types.trajectory import Trajectory
+
+
+class GoogleADKTraceProcessor:
+    """
+    Trace processor for Google ADK agents.
+
+    Captures events from the Google ADK runner and writes them to a JSONL file
+    for later conversion to standard trajectory format.
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+        self._lock = threading.Lock()
+
+    def _iso(self, ts: Optional[datetime] = None) -> str:
+        """Convert timestamp to ISO format"""
+        if ts is None:
+            ts = datetime.now(timezone.utc)
+        try:
+            return ts.isoformat()
+        except:
+            return datetime.now(timezone.utc).isoformat()
+
+    def _write(self, obj: Dict[str, Any]):
+        """Thread-safe write to JSONL 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, trace_id: str, metadata: Dict[str, Any]):
+        """Record trace start"""
+        self._write({
+            "record": "trace_start",
+            "trace_id": trace_id,
+            "metadata": metadata,
+            "started_at": self._iso(),
+            "ts": self._iso(),
+        })
+
+    def end_trace(self, trace_id: str):
+        """Record trace end"""
+        self._write({
+            "record": "trace_end",
+            "trace_id": trace_id,
+            "ended_at": self._iso(),
+            "ts": self._iso(),
+        })
+
+    def record_user_input(self, trace_id: str, user_input: str):
+        """Record user input"""
+        self._write({
+            "record": "user_input",
+            "trace_id": trace_id,
+            "content": user_input,
+            "ts": self._iso(),
+        })
+
+    def record_event(self, trace_id: str, event: Any):
+        """Record a Google ADK event"""
+        event_data = self._extract_event_data(event)
+        self._write({
+            "record": "event",
+            "trace_id": trace_id,
+            "event_type": type(event).__name__,
+            "event_data": event_data,
+            "ts": self._iso(),
+        })
+
+    def record_error(self, trace_id: str, error: str):
+        """Record an error"""
+        self._write({
+            "record": "error",
+            "trace_id": trace_id,
+            "error": error,
+            "ts": self._iso(),
+        })
+
+    def _extract_event_data(self, event: Any) -> Dict[str, Any]:
+        """Extract relevant data from a Google ADK event per ADK Events documentation"""
+        data = {}
+
+        # === Core Event Identifiers ===
+        if hasattr(event, 'id') and event.id:
+            data['event_id'] = str(event.id)
+
+        if hasattr(event, 'invocation_id') and event.invocation_id:
+            data['invocation_id'] = str(event.invocation_id)
+
+        if hasattr(event, 'timestamp') and event.timestamp:
+            data['event_timestamp'] = event.timestamp
+
+        # === Event Origin ===
+        if hasattr(event, 'author'):
+            data['author'] = str(event.author) if event.author else None
+
+        if hasattr(event, 'branch') and event.branch:
+            data['branch'] = str(event.branch)
+
+        # === Streaming/Completion Status ===
+        if hasattr(event, 'partial'):
+            data['partial'] = event.partial
+
+        if hasattr(event, 'turn_complete'):
+            data['turn_complete'] = event.turn_complete
+
+        # === Final Response Check (method, not attribute) ===
+        if hasattr(event, 'is_final_response') and callable(event.is_final_response):
+            try:
+                data['is_final_response'] = event.is_final_response()
+            except:
+                data['is_final_response'] = None
+
+        # === Content Extraction ===
+        if hasattr(event, 'content') and event.content:
+            content_data = self._extract_content(event.content)
+            if content_data:
+                data['content'] = content_data
+
+        # === Function Calls/Responses using helper methods ===
+        if hasattr(event, 'get_function_calls') and callable(event.get_function_calls):
+            try:
+                function_calls = event.get_function_calls()
+                if function_calls:
+                    data['function_calls'] = [
+                        {'name': fc.name, 'args': dict(fc.args) if fc.args else {}}
+                        for fc in function_calls
+                    ]
+            except:
+                pass
+
+        if hasattr(event, 'get_function_responses') and callable(event.get_function_responses):
+            try:
+                function_responses = event.get_function_responses()
+                if function_responses:
+                    data['function_responses'] = [
+                        {'name': fr.name, 'response': fr.response if hasattr(fr, 'response') else None}
+                        for fr in function_responses
+                    ]
+            except:
+                pass
+
+        # === Actions (EventActions object, not iterable) ===
+        if hasattr(event, 'actions') and event.actions:
+            data['actions'] = self._extract_actions(event.actions)
+
+        # === Error Fields ===
+        if hasattr(event, 'error_code') and event.error_code:
+            data['error_code'] = str(event.error_code)
+
+        if hasattr(event, 'error_message') and event.error_message:
+            data['error_message'] = str(event.error_message)
+
+        # === Long Running Tool IDs ===
+        if hasattr(event, 'long_running_tool_ids') and event.long_running_tool_ids:
+            data['long_running_tool_ids'] = list(event.long_running_tool_ids)
+
+        # === Legacy fields (keep for backwards compatibility) ===
+        if hasattr(event, 'tool_calls') and event.tool_calls:
+            data['tool_calls'] = self._extract_tool_calls(event.tool_calls)
+
+        if hasattr(event, 'tool_results') and event.tool_results:
+            data['tool_results'] = self._extract_tool_results(event.tool_results)
+
+        return data
+
+    def _extract_content(self, content: Any) -> Optional[Dict[str, Any]]:
+        """Extract content from Google ADK content object"""
+        if content is None:
+            return None
+
+        result = {}
+
+        if hasattr(content, 'role'):
+            result['role'] = str(content.role)
+
+        if hasattr(content, 'parts') and content.parts:
+            parts = []
+            for part in content.parts:
+                part_data = {}
+                if hasattr(part, 'text') and part.text:
+                    part_data['text'] = part.text
+                if hasattr(part, 'function_call') and part.function_call:
+                    part_data['function_call'] = self._extract_function_call(part.function_call)
+                if hasattr(part, 'function_response') and part.function_response:
+                    part_data['function_response'] = self._extract_function_response(part.function_response)
+                if part_data:
+                    parts.append(part_data)
+            if parts:
+                result['parts'] = parts
+
+        return result if result else None
+
+    def _extract_function_call(self, fc: Any) -> Dict[str, Any]:
+        """Extract function call data"""
+        data = {}
+        if hasattr(fc, 'name'):
+            data['name'] = fc.name
+        if hasattr(fc, 'args'):
+            try:
+                data['args'] = dict(fc.args) if fc.args else {}
+            except:
+                data['args'] = str(fc.args)
+        if hasattr(fc, 'id'):
+            data['id'] = fc.id
+        return data
+
+    def _extract_function_response(self, fr: Any) -> Dict[str, Any]:
+        """Extract function response data"""
+        data = {}
+        if hasattr(fr, 'name'):
+            data['name'] = fr.name
+        if hasattr(fr, 'response'):
+            try:
+                data['response'] = dict(fr.response) if isinstance(fr.response, dict) else str(fr.response)
+            except:
+                data['response'] = str(fr.response)
+        if hasattr(fr, 'id'):
+            data['id'] = fr.id
+        return data
+
+    def _extract_actions(self, actions: Any) -> Dict[str, Any]:
+        """
+        Extract actions data from EventActions object.
+
+        EventActions has these properties:
+        - state_delta: dict of state changes
+        - artifact_delta: dict of artifact changes
+        - transfer_to_agent: string, agent to transfer control to
+        - escalate: bool, signal to terminate loop
+        - skip_summarization: bool, skip LLM summarization of tool result
+        """
+        result = {}
+        try:
+            # State changes (dict of {key: value} pairs)
+            if hasattr(actions, 'state_delta') and actions.state_delta:
+                result['state_delta'] = dict(actions.state_delta)
+
+            # Artifact changes (dict of {filename: version})
+            if hasattr(actions, 'artifact_delta') and actions.artifact_delta:
+                result['artifact_delta'] = dict(actions.artifact_delta)
+
+            # Control flow: agent transfer
+            if hasattr(actions, 'transfer_to_agent') and actions.transfer_to_agent:
+                result['transfer_to_agent'] = str(actions.transfer_to_agent)
+
+            # Control flow: loop escalation/termination
+            if hasattr(actions, 'escalate') and actions.escalate:
+                result['escalate'] = bool(actions.escalate)
+
+            # Skip summarization flag for tool results
+            if hasattr(actions, 'skip_summarization') and actions.skip_summarization:
+                result['skip_summarization'] = bool(actions.skip_summarization)
+
+        except Exception as e:
+            result['_extraction_error'] = str(e)
+            result['_raw'] = str(actions)
+
+        return result
+
+    def _extract_tool_calls(self, tool_calls: Any) -> List[Dict[str, Any]]:
+        """Extract tool calls data"""
+        result = []
+        try:
+            for tc in tool_calls:
+                tc_data = {}
+                if hasattr(tc, 'name'):
+                    tc_data['name'] = tc.name
+                if hasattr(tc, 'args'):
+                    tc_data['args'] = dict(tc.args) if tc.args else {}
+                result.append(tc_data)
+        except:
+            result.append({'raw': str(tool_calls)})
+        return result
+
+    def _extract_tool_results(self, tool_results: Any) -> List[Dict[str, Any]]:
+        """Extract tool results data"""
+        result = []
+        try:
+            for tr in tool_results:
+                tr_data = {}
+                if hasattr(tr, 'name'):
+                    tr_data['name'] = tr.name
+                if hasattr(tr, 'result'):
+                    tr_data['result'] = tr.result
+                result.append(tr_data)
+        except:
+            result.append({'raw': str(tool_results)})
+        return result
+
+
+class GoogleADKTrajectoryConverter:
+    """Convert Google ADK traces to standard trajectory format"""
+
+    def __init__(self, output_dir: str, timestamp: Optional[str] = None):
+        self.output_dir = output_dir
+        self.timestamp = timestamp  # Used for output filename
+        os.makedirs(self.output_dir, exist_ok=True)
+
+    def load_traces(self, trace_file: str):
+        """Load traces from JSONL file
+
+        Returns:
+            traces_by_id: Dict mapping trace_id to list of all records (user_input, event, error) in order
+            trace_metadata: Dict mapping trace_id to metadata from trace_start
+            user_inputs: Dict mapping trace_id to list of user input strings (for backward compat)
+        """
+        if not os.path.exists(trace_file):
+            raise FileNotFoundError(f"Trace file not found: {trace_file}")
+
+        traces_by_id = defaultdict(list)
+        trace_metadata = {}
+        user_inputs = defaultdict(list)
+
+        with open(trace_file, 'r', encoding='utf-8') as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                data = json.loads(line)
+                trace_id = data.get("trace_id")
+
+                if data.get("record") == "trace_start":
+                    trace_metadata[trace_id] = data.get("metadata", {})
+
+                elif data.get("record") == "user_input":
+                    # Store in both places: ordered list for interleaving, separate list for backward compat
+                    user_inputs[trace_id].append(data.get("content", ""))
+                    traces_by_id[trace_id].append(data)
+
+                elif data.get("record") == "event":
+                    traces_by_id[trace_id].append(data)
+
+                elif data.get("record") == "error":
+                    traces_by_id[trace_id].append(data)
+
+        return traces_by_id, trace_metadata, user_inputs
+
+    def _parse_tool_input(self, input_data: Any) -> Dict[str, Any]:
+        """Parse tool input to dict"""
+        try:
+            if isinstance(input_data, dict):
+                return input_data
+            elif isinstance(input_data, str):
+                return json.loads(input_data)
+            else:
+                return {"raw": str(input_data)}
+        except:
+            return {"raw": str(input_data)}
+
+    def _parse_tool_output(self, output_data: Any) -> Any:
+        """
+        Parse tool output to extract the actual result data.
+
+        Google ADK MCP tool responses have a nested structure:
+        {
+            'content': [{'type': 'text', 'text': '{"data": ...}'}],
+            'isError': False,
+            'structuredContent': {'data': {...}}  # Already parsed
+        }
+
+        This method extracts the actual data to align with other SDK formats.
+        """
+        if output_data is None:
+            return None
+
+        # If it's a string, try to parse as JSON
+        if isinstance(output_data, str):
+            try:
+                return json.loads(output_data)
+            except (json.JSONDecodeError, TypeError):
+                return output_data
+
+        # If it's not a dict, return as-is
+        if not isinstance(output_data, dict):
+            return output_data
+
+        # Check for Google ADK MCP response structure
+        # Prefer structuredContent as it's already parsed
+        if 'structuredContent' in output_data and output_data['structuredContent']:
+            return output_data['structuredContent']
+
+        # Fall back to parsing content[0].text if available
+        if 'content' in output_data and isinstance(output_data['content'], list):
+            for item in output_data['content']:
+                if isinstance(item, dict) and item.get('type') == 'text':
+                    text = item.get('text', '')
+                    if text:
+                        try:
+                            return json.loads(text)
+                        except (json.JSONDecodeError, TypeError):
+                            return text
+
+        # If it has 'data' key directly, return that (already normalized)
+        if 'data' in output_data and len(output_data) == 1:
+            return output_data
+
+        # Return as-is if no special structure detected
+        return output_data
+
+    def _build_trajectory_steps(
+        self,
+        traj: Trajectory,
+        records: List[Dict],
+        tool_to_server: Optional[Dict[str, str]] = None
+    ):
+        """Build trajectory steps from Google ADK records (user_input, event, error)
+
+        Args:
+            traj: Trajectory object to append steps to
+            records: List of all records (user_input, event, error) in chronological order
+            tool_to_server: Optional mapping of tool_name -> server_name from config
+
+        Note: Records are processed in order to properly interleave user inputs
+        with agent responses for multi-turn conversations.
+        """
+        tool_to_server = tool_to_server or {}
+
+        # Process all records in chronological order 
+        for record in records:
+            record_type = record.get("record")
+
+            # Handle user input records
+            if record_type == "user_input":
+                user_input = record.get("content", "")
+                traj.append_user_step(user_input)
+                continue
+
+            # Handle error records
+            if record_type == "error":
+                traj.append_agent_step(
+                    action="error",
+                    metadata={"error": record.get("error", "Unknown error")}
+                )
+                continue
+
+            # Handle event records (record_type == "event")
+            event_data = record.get("event_data", {})
+            content = event_data.get("content", {})
+            is_final = event_data.get("is_final_response", False)
+
+            # Process function calls (tool invocations)
+            if content and "parts" in content:
+                for part in content.get("parts", []):
+                    # Function call (agent action)
+                    if "function_call" in part:
+                        fc = part["function_call"]
+                        tool_name = fc.get("name", "unknown_tool")
+                        tool_params = fc.get("args", {})
+
+                        # Look up server name from mapping, fallback to "unknown"
+                        server_name = tool_to_server.get(tool_name, "unknown")
+
+                        # Build action string
+                        action_str = f"{tool_name}("
+                        params = []
+                        for k, v in tool_params.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_params,
+                            server=server_name
+                        )
+
+                    # Function response (tool return)
+                    if "function_response" in part:
+                        fr = part["function_response"]
+                        tool_name = fr.get("name", "unknown_tool")
+                        result = fr.get("response", "")
+
+                        # Look up server name from mapping
+                        server_name = tool_to_server.get(tool_name, "unknown")
+
+                        parsed_result = self._parse_tool_output(result)
+
+                        traj.append_tool_return(
+                            result=parsed_result,
+                            tool_name=tool_name,
+                            server=server_name
+                        )
+
+                    # Final text response to user after each turn
+                    if is_final and "text" in part:
+                        traj.append_agent_step(
+                            action="send_message_to_user",
+                            metadata={"message": part["text"]}
+                        )
+                        # Set final response in trajectory
+                        traj.set_final_response(part["text"])
+
+    def convert_trace_to_trajectory(
+        self,
+        trace_id: str,
+        records: List[Dict],
+        metadata: Dict,
+        user_inputs: List[str]
+    ) -> Trajectory:
+        """Convert a single trace to trajectory format
+
+        Args:
+            trace_id: The trace ID
+            records: List of all records (user_input, event, error) in chronological order
+            metadata: Trace metadata from trace_start
+            user_inputs: List of user input strings (for instruction extraction)
+
+        Returns:
+            Trajectory object
+        """
+        task_id = metadata.get("task_id", trace_id[:16])
+        instruction = metadata.get("instruction", "")
+        if not instruction and user_inputs:
+            instruction = user_inputs[0]
+        malicious_instruction = metadata.get("malicious_goal")
+
+        # Get tool-to-server mapping from trace metadata
+        tool_to_server = metadata.get("tool_to_server", {})
+
+        # Create Trajectory instance
+        traj = Trajectory(
+            task_id=task_id,
+            original_instruction=instruction,
+            malicious_instruction=malicious_instruction,
+            domain=metadata.get("domain"),
+            risk_category=metadata.get("category")
+        )
+
+        # Calculate duration from records
+        if records:
+            timestamps = [r.get("ts") for r in records if r.get("ts")]
+            if len(timestamps) >= 2:
+                try:
+                    start = datetime.fromisoformat(timestamps[0].replace('Z', '+00:00'))
+                    end = datetime.fromisoformat(timestamps[-1].replace('Z', '+00:00'))
+                    duration = (end - start).total_seconds()
+                except:
+                    duration = 0.0
+            else:
+                duration = 0.0
+        else:
+            duration = 0.0
+
+        traj.data["traj_info"]["duration"] = round(duration, 3)
+        traj.data["traj_info"]["metadata"] = {"trace_id": trace_id}
+
+        # Build trajectory steps from records (properly interleaved for multi-turn)
+        self._build_trajectory_steps(traj, records, tool_to_server)
+
+        return traj
+
+    def process_trace_file(self, trace_file: str, output_name: Optional[str] = None) -> List[Trajectory]:
+        """Process trace file and generate trajectory files
+
+        Args:
+            trace_file: Path to the trace JSONL file
+            output_name: Optional name prefix for output files
+
+        Returns:
+            List of Trajectory objects
+        """
+        traces_by_id, trace_metadata, user_inputs = self.load_traces(trace_file)
+
+        trajectories = []
+
+        for i, (trace_id, records) in enumerate(traces_by_id.items()):
+            metadata = trace_metadata.get(trace_id, {})
+            task_id = metadata.get("task_id", trace_id[:16])
+            inputs = user_inputs.get(trace_id, [])
+
+            traj = self.convert_trace_to_trajectory(
+                trace_id, records, metadata, inputs
+            )
+
+            if self.timestamp:
+                if len(traces_by_id) == 1:
+                    filename = f"{self.timestamp}.json"
+                else:
+                    filename = f"{self.timestamp}_{i}.json"
+            else:
+                filename = f"{task_id}_{trace_id}.json"
+            output_path = os.path.join(self.output_dir, filename)
+
+            with open(output_path, 'w', encoding='utf-8') as f:
+                json.dump(traj.to_dict(), f, indent=2, ensure_ascii=False)
+
+            trajectories.append(traj)
+
+        return trajectories

agent/langchain/__init__.py

@@ -0,0 +1,8 @@
+"""
+LangChain Agent implementation for DecodingTrust Agent Arena.
+"""
+
+from .src.agent import LangChainAgent
+
+__all__ = ["LangChainAgent"]
+