1 line
15 KiB
JSON
1 line
15 KiB
JSON
{"repo": "aws-solutions-library-samples/accelerated-intelligent-document-processing-on-aws", "files": {"lib/idp_common_pkg/idp_common/monitoring/stepfunctions_service.py": {"3fa43cb1b3faf5c4b29e266f214bd1914eac3b0f": {"t": 1785846911.4556365, "n": 13999, "h": "93d1a01c45afd6ff", "c": "# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n# SPDX-License-Identifier: MIT-0\n\n\"\"\"\nStep Functions service for IDP execution analysis.\n\nProvides reusable functions to retrieve and analyse the processing history\nof a document through AWS Step Functions, including a timeline of every step\nthe document went through and where it failed.\n\nThis module is a pure library — it contains no ``@tool`` decorators and makes\nno assumptions about agent frameworks. Agent tool wrappers that call these\nfunctions live in ``agents/error_analyzer/tools/stepfunction_tool.py``.\n\nUsage::\n\n from idp_common.monitoring.stepfunctions_service import (\n get_execution_arn_from_document,\n get_execution_data,\n analyze_execution_timeline,\n extract_failure_details,\n )\n\n arn = get_execution_arn_from_document(document_record)\n data = get_execution_data(arn)\n timeline = analyze_execution_timeline(arn)\n failure = extract_failure_details(data[\"events\"])\n\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nfrom datetime import datetime, timezone\nfrom typing import Any, Dict, List, Optional, Set\n\nimport boto3\n\nlogger = logging.getLogger(__name__)\n\n# ---------------------------------------------------------------------------\n# Module-level lazy boto3 client cache, keyed by region so that callers passing\n# an explicit region always get the correct client, even if a default-region\n# client was already initialised first.\n# ---------------------------------------------------------------------------\n_sf_clients: Dict[Optional[str], Any] = {}\n\n\ndef _get_sf_client(region: Optional[str] = None) -> Any:\n \"\"\"Return (and lazily create) a per-region Step Functions boto3 client.\"\"\"\n if region not in _sf_clients:\n _sf_clients[region] = boto3.client(\"stepfunctions\", region_name=region)\n return _sf_clients[region]\n\n\n# ---------------------------------------------------------------------------\n# Event types that indicate a failure in the execution\n# ---------------------------------------------------------------------------\n_FAILURE_EVENT_TYPES: Set[str] = {\n \"TaskFailed\",\n \"TaskTimedOut\",\n \"ExecutionFailed\",\n \"ActivityFailed\",\n \"LambdaFunctionFailed\",\n \"LambdaFunctionTimedOut\",\n}\n\n# Mapping from event type to the key that holds failure detail in the event dict\n_FAILURE_DETAIL_KEYS: Dict[str, str] = {\n \"TaskFailed\": \"taskFailedEventDetails\",\n \"TaskTimedOut\": \"taskTimedOutEventDetails\",\n \"ExecutionFailed\": \"executionFailedEventDetails\",\n \"ActivityFailed\": \"activityFailedEventDetails\",\n \"LambdaFunctionFailed\": \"lambdaFunctionFailedEventDetails\",\n \"LambdaFunctionTimedOut\": \"lambdaFunctionTimedOutEventDetails\",\n}\n\n\n# ---------------------------------------------------------------------------\n# Public API\n# ---------------------------------------------------------------------------\n\n\ndef get_execution_arn_from_document(\n doc_record: Any,\n) -> str:\n \"\"\"\n Extract the Step Functions execution ARN from a DynamoDB document record.\n\n Accepts either a plain ``dict`` (raw DynamoDB item) or a\n :class:`~idp_common.monitoring.models.DocumentRecord` dataclass.\n\n Args:\n doc_record: DynamoDB document record dict or ``DocumentRecord``.\n\n Returns:\n Execution ARN string, or ``\"\"`` if not available.\n \"\"\"\n if isinstance(doc_record, dict):\n return (\n doc_record.get(\"WorkflowExecutionArn\", \"\")\n or doc_record.get(\"ExecutionArn\", \"\")\n or doc_record.get(\"workflow_execution_arn\", \"\")\n or \"\"\n )\n # Handle DocumentRecord dataclass (or any object with the attribute)\n return (\n getattr(doc_record, \"workflow_execution_arn\", \"\")\n or getattr(doc_record, \"WorkflowExecutionArn\", \"\")\n or \"\"\n )\n\n\ndef get_execution_data(\n execution_arn: str,\n region: Optional[str] = None,\n) -> Dict[str, Any]:\n \"\"\"\n Fetch full Step Functions execution history for a given execution ARN.\n\n Paginates through the ``get_execution_history`` API to retrieve all events.\n\n Args:\n execution_arn: Full ARN of the Step Functions execution.\n region: AWS region name. Defaults to the region inferred by\n boto3 from the environment.\n\n Returns:\n ``{\n \"execution_arn\": str,\n \"status\": str, # SUCCEEDED | FAILED | RUNNING | ABORTED | TIMED_OUT\n \"start_date\": str, # ISO 8601 or \"\"\n \"stop_date\": str, # ISO 8601 or \"\"\n \"events\": list, # Raw execution history events (timestamps serialised)\n \"input\": str, # Execution input JSON string\n }``\n \"\"\"\n sf = _get_sf_client(region)\n result: Dict[str, Any] = {\n \"execution_arn\": execution_arn,\n \"status\": \"UNKNOWN\",\n \"start_date\": \"\",\n \"stop_date\": \"\",\n \"events\": [],\n \"input\": \"\",\n }\n\n try:\n desc = sf.describe_execution(executionArn=execution_arn)\n result[\"status\"] = desc.get(\"status\", \"UNKNOWN\")\n\n start = desc.get(\"startDate\")\n stop = desc.get(\"stopDate\")\n result[\"start_date\"] = _serialize_datetime(start)\n result[\"stop_date\"] = _serialize_datetime(stop)\n result[\"input\"] = desc.get(\"input\", \"\")\n\n # Paginate execution history\n paginator = sf.get_paginator(\"get_execution_history\")\n events: List[Dict[str, Any]] = []\n for page in paginator.paginate(\n executionArn=execution_arn,\n includeExecutionData=True,\n ):\n events.extend(page.get(\"events\", []))\n\n # Serialise all timestamp fields so the result is JSON-safe\n for event in events:\n ts = event.get(\"timestamp\")\n if ts is not None and not isinstance(ts, str):\n event[\"timestamp\"] = _serialize_datetime(ts)\n\n result[\"events\"] = events\n\n except sf.exceptions.ExecutionDoesNotExist:\n # Use a distinct status so callers can tell \"not found\" from a\n # transient API error (which leaves status as \"UNKNOWN\").\n logger.warning(\"Execution not found: %s\", execution_arn)\n result[\"status\"] = \"NOT_FOUND\"\n except Exception as exc: # noqa: BLE001\n logger.warning(\"Failed to get execution data for %s: %s\", execution_arn, exc)\n\n return result\n\n\ndef analyze_execution_timeline(\n execution_arn: str,\n max_events: int = 200,\n region: Optional[str] = None,\n) -> Dict[str, Any]:\n \"\"\"\n Analyse a Step Functions execution and produce a structured timeline.\n\n Returns the sequence of states with their durations, identifying which\n state failed and extracting the error details.\n\n Args:\n execution_arn: Full ARN of the Step Functions execution.\n max_events: Maximum history events to process (default: 200).\n region: AWS region name.\n\n Returns:\n ``{\n \"execution_arn\": str,\n \"overall_status\": str,\n \"total_duration_ms\": float,\n \"states\": [\n {\n \"name\": str,\n \"status\": str, # SUCCEEDED | FAILED\n \"start_time\": str,\n \"end_time\": str,\n \"duration_ms\": float,\n \"is_failure\": bool,\n }\n ],\n \"failed_state\": str,\n \"failure_details\": dict,\n }``\n \"\"\"\n exec_data = get_execution_data(execution_arn, region=region)\n all_events: List[Dict[str, Any]] = exec_data.get(\"events\", [])\n if len(all_events) > max_events:\n logger.warning(\n \"Execution %s has %d events; truncating to %d for timeline analysis. \"\n \"Increase max_events to avoid missing failure details.\",\n execution_arn,\n len(all_events),\n max_events,\n )\n events = all_events[:max_events]\n\n timeline: Dict[str, Any] = {\n \"execution_arn\": execution_arn,\n \"overall_status\": exec_data.get(\"status\", \"UNKNOWN\"),\n \"total_duration_ms\": 0.0,\n \"states\": [],\n \"failed_state\": \"\",\n \"failure_details\": {},\n }\n\n state_starts: Dict[str, str] = {} # state_name → start_timestamp\n states: List[Dict[str, Any]] = []\n\n for event in events:\n event_type: str = event.get(\"type\", \"\")\n timestamp: str = event.get(\"timestamp\", \"\")\n\n if event_type == \"TaskStateEntered\":\n state_name = event.get(\"stateEnteredEventDetails\", {}).get(\"name\", \"\")\n if state_name:\n state_starts[state_name] = timestamp\n continue\n\n # --- State-completion and failure events that produce a timeline entry ---\n if event_type == \"TaskStateExited\":\n state_name = event.get(\"stateExitedEventDetails\", {}).get(\"name\", \"\")\n is_failure = False\n status = \"SUCCEEDED\"\n\n elif event_type in (\"TaskFailed\", \"TaskTimedOut\"):\n # Task-level failures do not carry stateExitedEventDetails — look\n # back to the most recently entered state to get the name.\n state_name = list(state_starts.keys())[-1] if state_starts else \"Unknown\"\n is_failure = True\n status = \"FAILED\"\n\n elif event_type == \"ExecutionFailed\":\n state_name = event.get(\"stateExitedEventDetails\", {}).get(\"name\", \"\")\n if not state_name:\n # Attribute failure to the last known entered state\n state_name = (\n list(state_starts.keys())[-1] if state_starts else \"Unknown\"\n )\n is_failure = True\n status = \"FAILED\"\n\n else:\n # Ignore all other event types (ExecutionStarted, LambdaScheduled, etc.)\n continue\n\n start = state_starts.get(state_name, \"\")\n duration_ms = _compute_duration_ms(start, timestamp)\n\n states.append(\n {\n \"name\": state_name,\n \"status\": status,\n \"start_time\": start,\n \"end_time\": timestamp,\n \"duration_ms\": round(duration_ms, 1),\n \"is_failure\": is_failure,\n }\n )\n\n if is_failure and not timeline[\"failed_state\"]:\n timeline[\"failed_state\"] = state_name\n\n timeline[\"states\"] = states\n timeline[\"failure_details\"] = extract_failure_details(events)\n\n # Total execution duration\n start_date = exec_data.get(\"start_date\", \"\")\n stop_date = exec_data.get(\"stop_date\", \"\")\n if start_date and stop_date:\n timeline[\"total_duration_ms\"] = round(\n _compute_duration_ms(start_date, stop_date), 1\n )\n\n return timeline\n\n\ndef extract_failure_details(events: List[Dict[str, Any]]) -> Dict[str, Any]:\n \"\"\"\n Parse Step Functions execution history events to extract failure information.\n\n Handles all six failure event types:\n ``TaskFailed``, ``TaskTimedOut``, ``ExecutionFailed``,\n ``ActivityFailed``, ``LambdaFunctionFailed``, ``LambdaFunctionTimedOut``.\n\n Args:\n events: List of raw execution history event dicts (timestamps may be\n strings or ``datetime`` objects).\n\n Returns:\n ``{\n \"error\": str, # Error type (e.g. \"ThrottlingException\")\n \"cause\": str, # Error cause message (may be a JSON string)\n \"failed_state\": str, # Name of the state that failed\n \"event_type\": str, # The Step Functions event type\n }``\n All fields are empty strings if no failure event is found.\n \"\"\"\n result: Dict[str, Any] = {\n \"error\": \"\",\n \"cause\": \"\",\n \"failed_state\": \"\",\n \"event_type\": \"\",\n }\n\n # Walk in reverse to find the most recent failure event\n for event in reversed(events):\n event_type: str = event.get(\"type\", \"\")\n if event_type not in _FAILURE_EVENT_TYPES:\n continue\n\n result[\"event_type\"] = event_type\n\n detail_key = _FAILURE_DETAIL_KEYS.get(event_type, \"\")\n details: Dict[str, Any] = event.get(detail_key, {}) if detail_key else {}\n\n result[\"error\"] = details.get(\"error\", \"\") or details.get(\"Error\", \"\") or \"\"\n result[\"cause\"] = details.get(\"cause\", \"\") or details.get(\"Cause\", \"\") or \"\"\n\n # Trace back to find which state was entered before this failure\n event_id: int = event.get(\"id\", 0)\n for prior in reversed(events):\n if (\n prior.get(\"id\", 0) < event_id\n and prior.get(\"type\") == \"TaskStateEntered\"\n ):\n result[\"failed_state\"] = prior.get(\"stateEnteredEventDetails\", {}).get(\n \"name\", \"\"\n )\n break\n\n break # Use the most recent failure only\n\n return result\n\n\n# ---------------------------------------------------------------------------\n# Private helpers\n# ---------------------------------------------------------------------------\n\n\ndef _serialize_datetime(dt: Any) -> str:\n \"\"\"Convert a datetime object (or None) to an ISO 8601 string.\"\"\"\n if dt is None:\n return \"\"\n if isinstance(dt, str):\n return dt\n if isinstance(dt, datetime):\n if dt.tzinfo is None:\n dt = dt.replace(tzinfo=timezone.utc)\n return dt.isoformat()\n return str(dt)\n\n\ndef _compute_duration_ms(start: str, end: str) -> float:\n \"\"\"\n Compute the duration in milliseconds between two ISO 8601 timestamp strings.\n\n Returns 0.0 if either string is empty or unparseable.\n \"\"\"\n if not start or not end:\n return 0.0\n try:\n t0 = datetime.fromisoformat(start.replace(\"Z\", \"+00:00\"))\n t1 = datetime.fromisoformat(end.replace(\"Z\", \"+00:00\"))\n return (t1 - t0).total_seconds() * 1000\n except (ValueError, TypeError) as exc:\n logger.debug(\n \"Could not compute duration between '%s' and '%s': %s\", start, end, exc\n )\n return 0.0\n"}}}} |