aip-agents-binary 0.6.4__py3-none-any.whl

aip_agents/utils/sse_chunk_transformer.pyi

@@ -0,0 +1,166 @@
+from _typeshed import Incomplete
+from aip_agents.schema.a2a import A2AEvent
+from collections.abc import AsyncGenerator
+from enum import StrEnum
+from typing import Any
+
+__all__ = ['SSEChunkTransformer', 'TaskState', 'ChunkStatus', 'ChunkReason', 'ChunkFieldKeys']
+
+class TaskState(StrEnum):
+    """Task state values for SSE chunks."""
+    WORKING: str
+    COMPLETED: str
+    FAILED: str
+
+class ChunkStatus(StrEnum):
+    """Status values for SSE chunks."""
+    SUCCESS: str
+    ERROR: str
+
+class ChunkReason(StrEnum):
+    """Reason codes for special chunk states."""
+    EMPTY_PAYLOAD: str
+
+class ChunkFieldKeys(StrEnum):
+    """Field name constants for SSE chunk structure."""
+    STATUS: str
+    TASK_STATE: str
+    CONTENT: str
+    EVENT_TYPE: str
+    FINAL: str
+    METADATA: str
+    TIMESTAMP: str
+    TASK_ID: str
+    CONTEXT_ID: str
+    ARTIFACTS: str
+    REASON: str
+
+class SSEChunkTransformer:
+    '''Transforms A2AEvent stream to SSE-compatible output.
+
+    This class converts events from arun_a2a_stream into the normalized dict format
+    matching A2AConnector.astream_to_agent output.
+
+    Lifecycle:
+        Single-stream instance. Must NOT be reused across concurrent streams.
+        Each arun_sse_stream call creates a fresh instance.
+
+    Attributes:
+        task_id: Optional task identifier for the stream.
+        context_id: Optional context identifier for the stream.
+
+    Example:
+        >>> transformer = SSEChunkTransformer(task_id="task-123")
+        >>> async for chunk in transformer.transform_stream(agent.arun_a2a_stream("query")):
+        ...     print(chunk)
+    '''
+    task_id: Incomplete
+    context_id: Incomplete
+    def __init__(self, task_id: str | None = None, context_id: str | None = None, pii_mapping: dict[str, str] | None = None) -> None:
+        """Initialize the transformer with optional task and context IDs.
+
+        Args:
+            task_id: Optional task identifier for the stream.
+            context_id: Optional context identifier for the stream.
+            pii_mapping: Optional PII mapping to inject into each chunk's metadata.
+        """
+    @staticmethod
+    def normalize_metadata_enums(data: Any) -> Any:
+        """Recursively convert enum keys/values to their string values.
+
+        This is a pure normalization utility that converts any enum instances
+        (MetadataFieldKeys, Kind, Status, etc.) to their .value strings.
+
+        Args:
+            data: Dict, list, or value that may contain enum keys/values.
+
+        Returns:
+            Normalized data with all enums converted to their .value strings.
+        """
+    @staticmethod
+    def normalize_event_type_value(event_type: Any) -> str | None:
+        """Convert A2AStreamEventType enum to string.
+
+        Args:
+            event_type: Event type (enum, string, or None).
+
+        Returns:
+            String value of the event type, or None if invalid.
+        """
+    @staticmethod
+    def create_artifact_hash(artifact: dict[str, Any]) -> str:
+        """Create a stable hash for artifact deduplication.
+
+        Uses name, content_type, mime_type, and file_data for hashing,
+        excluding artifact_id which may be randomly generated.
+
+        Args:
+            artifact: Artifact dict with name, content_type, mime_type, and optionally file_data.
+
+        Returns:
+            SHA256 hexdigest hash string for deduplication.
+        """
+    @staticmethod
+    def extract_tool_outputs(tool_calls: list[dict[str, Any]]) -> list[str]:
+        """Extract human-readable output strings from tool calls.
+
+        Args:
+            tool_calls: List of tool call dictionaries.
+
+        Returns:
+            List of human-readable output strings.
+        """
+    @staticmethod
+    def format_tool_output(output: Any, tool_name: str) -> str:
+        """Format a single tool output for display.
+
+        Args:
+            output: The tool output to format.
+            tool_name: The name of the tool.
+
+        Returns:
+            The formatted output string.
+        """
+    @staticmethod
+    def apply_hitl_content_override(content: str | None, event_type_str: str, metadata: dict[str, Any]) -> str | None:
+        """Apply HITL content override when HITL is active and tool results are available.
+
+        This method overrides the content with human-readable tool output when HITL
+        is active, matching A2AConnector behavior.
+
+        Args:
+            content: The original content/status message.
+            event_type_str: The type of event being processed (normalized string).
+            metadata: The metadata dictionary containing tool_info and hitl flag.
+
+        Returns:
+            The original content or human-readable tool output if HITL is active.
+        """
+    def transform_event(self, event: A2AEvent) -> dict[str, Any]:
+        """Transform a single A2AEvent to SSE chunk format.
+
+        Converts the A2AEvent structure to the normalized SSE chunk format,
+        relocating fields like tool_info and thinking_and_activity_info into
+        metadata, and normalizing enum values to strings.
+
+        Args:
+            event: Single A2AEvent dict from arun_a2a_stream.
+
+        Returns:
+            SSEChunk dict with normalized structure.
+        """
+    async def transform_stream(self, stream: AsyncGenerator[A2AEvent, None]) -> AsyncGenerator[dict[str, Any], None]:
+        """Transform A2AEvent stream to SSE-compatible chunks.
+
+        Wraps the input stream and transforms each event, handling artifact
+        deduplication and time tracking across the stream.
+
+        Args:
+            stream: Async generator yielding A2AEvent dicts.
+
+        Yields:
+            SSEChunk dicts with normalized structure.
+
+        Raises:
+            Exceptions from underlying stream propagate to caller.
+        """

aip_agents/utils/step_limit_manager.py

@@ -0,0 +1,265 @@
+"""Step limit manager for enforcing execution and delegation limits.
+
+Authors:
+    Saul Sayers (saul.sayers@gdplabs.id)
+"""
+
+from contextvars import ContextVar
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from aip_agents.schema.step_limit import (
+    MaxDelegationDepthExceededError,
+    MaxStepsExceededError,
+    StepLimitConfig,
+    StepLimitErrorResponse,
+    StepLimitErrorType,
+)
+
+# Context variables for cross-agent propagation
+_DELEGATION_DEPTH_CVAR: ContextVar[int] = ContextVar("delegation_depth", default=0)
+_REMAINING_STEP_BUDGET_CVAR: ContextVar[int | None] = ContextVar("remaining_step_budget", default=None)
+_DELEGATION_CHAIN_CVAR: ContextVar[tuple[str, ...]] = ContextVar("delegation_chain", default=())
+_STEP_LIMIT_CONFIG_CVAR: ContextVar[StepLimitConfig | None] = ContextVar("step_limit_config", default=None)
+
+
+@dataclass
+class StepExecutionContext:
+    """Runtime context for tracking step execution and delegation depth.
+
+    Attributes:
+        current_step: Current step number (0-indexed).
+        delegation_depth: Current depth in delegation chain (0 for root agent).
+        remaining_step_budget: Steps remaining before limit is hit.
+        delegation_chain: List of agent names in the delegation chain.
+    """
+
+    current_step: int = 0
+    delegation_depth: int = 0
+    remaining_step_budget: int | None = None
+    delegation_chain: list[str] = field(default_factory=list)
+
+
+class StepLimitManager:
+    """Manages step and delegation limit enforcement during agent execution.
+
+    This manager integrates with LangGraph's existing step mechanisms and adds
+    delegation depth tracking and budget propagation.
+
+    Attributes:
+        config: Step limit configuration.
+        context: Current execution context.
+    """
+
+    def __init__(
+        self,
+        config: StepLimitConfig | None = None,
+        initial_delegation_depth: int = 0,
+        parent_step_budget: int | None = None,
+    ) -> None:
+        """Initialize step limit manager.
+
+        Args:
+            config: Optional step limit configuration. Uses defaults if None.
+            initial_delegation_depth: Starting delegation depth (from parent).
+            parent_step_budget: Remaining step budget inherited from parent agent.
+        """
+        self.config = config or StepLimitConfig()
+        self.context = StepExecutionContext(
+            delegation_depth=initial_delegation_depth,
+            remaining_step_budget=parent_step_budget if parent_step_budget is not None else self.config.max_steps,
+        )
+
+    def check_step_limit(self, agent_name: str = "agent", count: int = 1) -> None:
+        """Check if taking 'count' steps would exceed limit.
+
+        Args:
+            agent_name: Name of the agent to identify in error message.
+            count: Number of steps to check (useful for parallel tool batches).
+
+        Raises:
+            MaxStepsExceededError: If max_steps limit is exceeded.
+        """
+        limit_exceeded = False
+        limit_val = self.config.max_steps
+        current_val = self.context.current_step
+
+        if self.context.current_step + count > self.config.max_steps:
+            limit_exceeded = True
+        elif self.context.remaining_step_budget is not None and self.context.remaining_step_budget < count:
+            limit_exceeded = True
+            # If we hit budget, the limit was effectively what we've done so far + what's left
+            limit_val = self.context.current_step + self.context.remaining_step_budget
+
+        if limit_exceeded:
+            error_response = StepLimitErrorResponse(
+                error_type=StepLimitErrorType.STEP_LIMIT_EXCEEDED,
+                agent_name=agent_name,
+                current_value=current_val,
+                configured_limit=limit_val,
+                message=f"Agent exceeded maximum step limit ({current_val + count}/{limit_val} steps)",
+            )
+            raise MaxStepsExceededError(error_response)
+
+    def check_delegation_depth(self, target_agent_name: str) -> None:
+        """Check if delegation to target agent would exceed depth limit.
+
+        Args:
+            target_agent_name: Name of the agent to delegate to.
+
+        Raises:
+            MaxDelegationDepthExceededError: If delegation depth limit exceeded.
+        """
+        # Check for circular delegation
+        if target_agent_name in self.context.delegation_chain:
+            error_response = StepLimitErrorResponse(
+                error_type=StepLimitErrorType.DELEGATION_DEPTH_EXCEEDED,
+                agent_name=target_agent_name,
+                current_value=self.context.delegation_depth,
+                configured_limit=self.config.max_delegation_depth,
+                message=f"Circular delegation detected: '{target_agent_name}' already in chain {self.context.delegation_chain}",
+                delegation_chain=self.context.delegation_chain,
+            )
+            raise MaxDelegationDepthExceededError(error_response)
+
+        # Check depth limit
+        if self.context.delegation_depth >= self.config.max_delegation_depth:
+            error_response = StepLimitErrorResponse(
+                error_type=StepLimitErrorType.DELEGATION_DEPTH_EXCEEDED,
+                agent_name=target_agent_name,
+                current_value=self.context.delegation_depth,
+                configured_limit=self.config.max_delegation_depth,
+                message=f"Cannot delegate to '{target_agent_name}': Maximum delegation depth ({self.context.delegation_depth}/{self.config.max_delegation_depth}) exceeded",
+                delegation_chain=self.context.delegation_chain,
+            )
+            raise MaxDelegationDepthExceededError(error_response)
+
+    def increment_step(self, count: int = 1) -> None:
+        """Increment step counter and update remaining budget.
+
+        Args:
+            count: Number of steps to consume (defaults to 1).
+        """
+        if count <= 0:
+            return
+        self.context.current_step += count
+        if self.context.remaining_step_budget is not None:
+            self.context.remaining_step_budget = max(0, self.context.remaining_step_budget - count)
+
+    def get_child_budget(self, child_max_steps: int | None = None) -> int:
+        """Calculate step budget to allocate to child agent.
+
+        Algorithm:
+        1. If remaining_step_budget is None (root with no limit), use config.max_steps - 1
+        2. If remaining_step_budget <= 1, return 0 (no budget left for child)
+        3. Calculate child_budget = remaining_step_budget - 1 (reserve 1 for parent)
+        4. If child has own max_steps config, return min(child_budget, child.max_steps)
+        5. Otherwise return child_budget
+
+        Args:
+            child_max_steps: Optional child agent's own max_steps limit.
+
+        Returns:
+            Step budget for child agent, accounting for parent's continuation.
+            Returns 0 if no budget available for child.
+
+        Edge Cases:
+            - remaining=1: Returns 0 (parent needs the last step)
+            - remaining=None: Uses config.max_steps - 1
+            - child has own limit: Returns min(calculated_budget, child_limit)
+        """
+        if self.context.remaining_step_budget is None:
+            child_budget = self.config.max_steps - 1
+        elif self.context.remaining_step_budget <= 1:
+            return 0
+        else:
+            child_budget = self.context.remaining_step_budget - 1
+
+        if child_max_steps is not None:
+            return min(child_budget, child_max_steps)
+        return child_budget
+
+    def add_to_delegation_chain(self, agent_name: str) -> None:
+        """Add agent to delegation chain for tracking.
+
+        Args:
+            agent_name: Name of the agent being delegated to.
+        """
+        self.context.delegation_chain.append(agent_name)
+
+    @classmethod
+    def from_state(
+        cls,
+        state: dict[str, Any],
+        config: StepLimitConfig | None = None,
+    ) -> "StepLimitManager":
+        """Create manager from LangGraph state.
+
+        Args:
+            state: LangGraph agent state containing remaining_steps, etc.
+            config: Optional step limit configuration.
+
+        Returns:
+            Initialized step limit manager.
+        """
+        raw_config = state.get("step_limit_config")
+        if isinstance(raw_config, dict):
+            step_limit_config = StepLimitConfig(**raw_config)
+        else:
+            step_limit_config = raw_config or config or StepLimitConfig()
+
+        delegation_depth = state.get("delegation_depth")
+        delegation_chain = state.get("delegation_chain")
+        current_step = state.get("current_step", 0)
+
+        # Restore remaining budget
+        remaining_step_budget = state.get("remaining_step_budget")
+
+        # Try to read from ContextVars if not in state (fallback for first step)
+        if remaining_step_budget is None:
+            try:
+                remaining_step_budget = _REMAINING_STEP_BUDGET_CVAR.get()
+            except LookupError:  # pragma: no cover - unreachable with default value, defensive code
+                pass
+
+        if delegation_depth is None:
+            try:
+                delegation_depth = _DELEGATION_DEPTH_CVAR.get()
+            except LookupError:  # pragma: no cover - unreachable with default value, defensive code
+                delegation_depth = 0
+
+        if delegation_chain is None:
+            try:
+                delegation_chain = list(_DELEGATION_CHAIN_CVAR.get())
+            except LookupError:  # pragma: no cover - unreachable with default value, defensive code
+                delegation_chain = []
+
+        manager = cls(
+            config=step_limit_config,
+            initial_delegation_depth=delegation_depth,
+            parent_step_budget=remaining_step_budget,
+        )
+        manager.context.delegation_chain = list(delegation_chain)
+        manager.context.current_step = current_step
+        return manager
+
+    def set_context(self) -> None:
+        """Set context variables for downstream consumption (e.g. by delegation tools)."""
+        _REMAINING_STEP_BUDGET_CVAR.set(self.context.remaining_step_budget)
+        _DELEGATION_DEPTH_CVAR.set(self.context.delegation_depth)
+        _DELEGATION_CHAIN_CVAR.set(tuple(self.context.delegation_chain))
+        _STEP_LIMIT_CONFIG_CVAR.set(self.config)
+
+    def to_state_update(self) -> dict[str, Any]:
+        """Convert current context to LangGraph state update.
+
+        Returns:
+            Dictionary of state fields to update.
+        """
+        return {
+            "current_step": self.context.current_step,
+            "delegation_depth": self.context.delegation_depth,
+            "delegation_chain": self.context.delegation_chain,
+            "step_limit_config": asdict(self.config),
+            "remaining_step_budget": self.context.remaining_step_budget,
+        }

aip_agents/utils/step_limit_manager.pyi

@@ -0,0 +1,112 @@
+from _typeshed import Incomplete
+from aip_agents.schema.step_limit import MaxDelegationDepthExceededError as MaxDelegationDepthExceededError, MaxStepsExceededError as MaxStepsExceededError, StepLimitConfig as StepLimitConfig, StepLimitErrorResponse as StepLimitErrorResponse, StepLimitErrorType as StepLimitErrorType
+from dataclasses import dataclass, field
+from typing import Any
+
+@dataclass
+class StepExecutionContext:
+    """Runtime context for tracking step execution and delegation depth.
+
+    Attributes:
+        current_step: Current step number (0-indexed).
+        delegation_depth: Current depth in delegation chain (0 for root agent).
+        remaining_step_budget: Steps remaining before limit is hit.
+        delegation_chain: List of agent names in the delegation chain.
+    """
+    current_step: int = ...
+    delegation_depth: int = ...
+    remaining_step_budget: int | None = ...
+    delegation_chain: list[str] = field(default_factory=list)
+
+class StepLimitManager:
+    """Manages step and delegation limit enforcement during agent execution.
+
+    This manager integrates with LangGraph's existing step mechanisms and adds
+    delegation depth tracking and budget propagation.
+
+    Attributes:
+        config: Step limit configuration.
+        context: Current execution context.
+    """
+    config: Incomplete
+    context: Incomplete
+    def __init__(self, config: StepLimitConfig | None = None, initial_delegation_depth: int = 0, parent_step_budget: int | None = None) -> None:
+        """Initialize step limit manager.
+
+        Args:
+            config: Optional step limit configuration. Uses defaults if None.
+            initial_delegation_depth: Starting delegation depth (from parent).
+            parent_step_budget: Remaining step budget inherited from parent agent.
+        """
+    def check_step_limit(self, agent_name: str = 'agent', count: int = 1) -> None:
+        """Check if taking 'count' steps would exceed limit.
+
+        Args:
+            agent_name: Name of the agent to identify in error message.
+            count: Number of steps to check (useful for parallel tool batches).
+
+        Raises:
+            MaxStepsExceededError: If max_steps limit is exceeded.
+        """
+    def check_delegation_depth(self, target_agent_name: str) -> None:
+        """Check if delegation to target agent would exceed depth limit.
+
+        Args:
+            target_agent_name: Name of the agent to delegate to.
+
+        Raises:
+            MaxDelegationDepthExceededError: If delegation depth limit exceeded.
+        """
+    def increment_step(self, count: int = 1) -> None:
+        """Increment step counter and update remaining budget.
+
+        Args:
+            count: Number of steps to consume (defaults to 1).
+        """
+    def get_child_budget(self, child_max_steps: int | None = None) -> int:
+        """Calculate step budget to allocate to child agent.
+
+        Algorithm:
+        1. If remaining_step_budget is None (root with no limit), use config.max_steps - 1
+        2. If remaining_step_budget <= 1, return 0 (no budget left for child)
+        3. Calculate child_budget = remaining_step_budget - 1 (reserve 1 for parent)
+        4. If child has own max_steps config, return min(child_budget, child.max_steps)
+        5. Otherwise return child_budget
+
+        Args:
+            child_max_steps: Optional child agent's own max_steps limit.
+
+        Returns:
+            Step budget for child agent, accounting for parent's continuation.
+            Returns 0 if no budget available for child.
+
+        Edge Cases:
+            - remaining=1: Returns 0 (parent needs the last step)
+            - remaining=None: Uses config.max_steps - 1
+            - child has own limit: Returns min(calculated_budget, child_limit)
+        """
+    def add_to_delegation_chain(self, agent_name: str) -> None:
+        """Add agent to delegation chain for tracking.
+
+        Args:
+            agent_name: Name of the agent being delegated to.
+        """
+    @classmethod
+    def from_state(cls, state: dict[str, Any], config: StepLimitConfig | None = None) -> StepLimitManager:
+        """Create manager from LangGraph state.
+
+        Args:
+            state: LangGraph agent state containing remaining_steps, etc.
+            config: Optional step limit configuration.
+
+        Returns:
+            Initialized step limit manager.
+        """
+    def set_context(self) -> None:
+        """Set context variables for downstream consumption (e.g. by delegation tools)."""
+    def to_state_update(self) -> dict[str, Any]:
+        """Convert current context to LangGraph state update.
+
+        Returns:
+            Dictionary of state fields to update.
+        """

aip_agents/utils/token_usage_helper.py

@@ -0,0 +1,156 @@
+"""This module provides utility functions for handling token usage in LangGraph agents.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from typing import Any
+
+from langchain_core.messages.ai import AIMessage, UsageMetadata
+from langchain_core.messages.tool import ToolMessage
+from langgraph.types import Command
+
+# Constants for token usage field names
+USAGE_METADATA_KEY = "usage_metadata"
+TOTAL_USAGE_KEY = "total_usage"
+STEP_USAGE_KEY = "step_usage"
+
+
+def add_usage_metadata(
+    cur_accumulated_token_usage: UsageMetadata | None, new_token_usage: UsageMetadata | None
+) -> UsageMetadata | None:
+    """Reducer function to accumulate UsageMetadata across agent runs.
+
+    Args:
+        cur_accumulated_token_usage: The current accumulated token usage metadata.
+        new_token_usage: New token usage metadata to add.
+
+    Returns:
+        Accumulated usage metadata or None if both inputs are None.
+    """
+    # Fast path for common cases
+    if cur_accumulated_token_usage is None:
+        return new_token_usage
+    if new_token_usage is None:
+        return cur_accumulated_token_usage
+
+    # Pre-allocate result dictionary with known size
+    keys = {"input_tokens", "output_tokens", "total_tokens"}
+    result = {key: cur_accumulated_token_usage.get(key, 0) + new_token_usage.get(key, 0) for key in keys}
+
+    # Only process details if they exist
+    detail_keys = {"input_token_details", "output_token_details"}
+    for key in detail_keys:
+        if cur_accumulated_token_usage.get(key) or new_token_usage.get(key):
+            result[key] = _merge_token_details(cur_accumulated_token_usage.get(key, {}), new_token_usage.get(key, {}))
+
+    return result
+
+
+def _merge_token_details(current: dict[str, int], new: dict[str, int]) -> dict[str, int]:
+    """Merge token details from two dictionaries.
+
+    Args:
+        current: Current token details.
+        new: New token details to add.
+
+    Returns:
+        Merged token details.
+    """
+    # Use dict comprehension for better performance
+    all_keys = current.keys() | new.keys()
+    return {key: current.get(key, 0) + new.get(key, 0) for key in all_keys}
+
+
+def extract_and_update_token_usage_from_ai_message(ai_message: AIMessage) -> dict[str, Any]:
+    """Extract token usage from AI message and prepare state update.
+
+    Args:
+        ai_message: The AI message containing usage metadata.
+
+    Returns:
+        Dictionary with accumulated_usage_metadata update if usage metadata is available.
+    """
+    state_updates = {}
+
+    if hasattr(ai_message, USAGE_METADATA_KEY) and ai_message.usage_metadata:
+        state_updates[TOTAL_USAGE_KEY] = ai_message.usage_metadata
+
+    return state_updates
+
+
+def extract_token_usage_from_tool_output(tool_output: Any) -> UsageMetadata | None:
+    """Extract token usage from various tool output formats.
+
+    Supports multiple tool output formats:
+    1. Dictionary with 'usage_metadata' field
+    2. Command with 'usage_metadata' attribute
+    3. Any object with 'usage_metadata' attribute
+
+    Args:
+        tool_output: The output from a tool execution.
+
+    Returns:
+        UsageMetadata if found, None otherwise.
+    """
+    # Handle dictionary format: {"result": "...", "usage_metadata": {...}}
+    if isinstance(tool_output, dict) and USAGE_METADATA_KEY in tool_output:
+        return tool_output[USAGE_METADATA_KEY]
+
+    # Handle Command format with usage_metadata in update field
+    if isinstance(tool_output, Command):
+        return extract_token_usage_from_command(tool_output)
+
+    # Handle any object with usage_metadata attribute
+    if hasattr(tool_output, USAGE_METADATA_KEY):
+        return getattr(tool_output, USAGE_METADATA_KEY)
+
+    return None
+
+
+def extract_token_usage_from_command(command: Command) -> UsageMetadata | None:
+    """Extract token usage from Command object.
+
+    Args:
+        command: The Command object to extract token usage from.
+
+    Returns:
+        UsageMetadata if found, None otherwise.
+    """
+    update_dict = getattr(command, "update", {}) or {}
+    if USAGE_METADATA_KEY in update_dict:
+        return update_dict[USAGE_METADATA_KEY]
+
+    # if not try to get from messages in update, and accumulate token usage
+    messages: list[ToolMessage] = update_dict.get("messages", [])
+    if not messages:
+        return None
+
+    accumulated_token_usage = None
+    for message in messages:
+        if USAGE_METADATA_KEY in message.response_metadata:
+            accumulated_token_usage = add_usage_metadata(
+                accumulated_token_usage, message.response_metadata[USAGE_METADATA_KEY]
+            )
+    return accumulated_token_usage
+
+
+def extract_token_usage_from_agent_response(agent_response: dict[str, Any]) -> UsageMetadata | None:
+    """Extract accumulated token usage from agent response.
+
+    Args:
+        agent_response: The agent response to extract token usage from.
+
+    Returns:
+        UsageMetadata if found, None otherwise.
+    """
+    if not isinstance(agent_response, dict):
+        return None
+
+    full_state = agent_response.get("full_final_state", {})
+    if not isinstance(full_state, dict):
+        return None
+
+    if TOTAL_USAGE_KEY in full_state:
+        return full_state[TOTAL_USAGE_KEY]
+    return None