litellm/litellm/a2a_protocol/utils.py

"""
Utility functions for A2A protocol.
"""

import hashlib
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final

import litellm
from litellm._logging import verbose_logger

if TYPE_CHECKING:
    from a2a.types import SendMessageRequest, SendStreamingMessageRequest


class A2ARequestUtils:
    """Utility class for A2A request/response processing."""

    @staticmethod
    def extract_text_from_message(message: Any) -> str:
        """
        Extract text content from A2A message parts.

        Args:
            message: A2A message dict or object with 'parts' containing text parts

        Returns:
            Concatenated text from all text parts
        """
        if message is None:
            return ""

        # Handle both dict and object access
        if isinstance(message, dict):
            parts = message.get("parts", [])
        else:
            parts = getattr(message, "parts", []) or []

        text_parts: Final[list[str]] = []
        for part in parts:
            if isinstance(part, dict):
                if part.get("kind") == "text":
                    text_parts.append(part.get("text", ""))
            else:
                if getattr(part, "kind", None) == "text":
                    text_parts.append(getattr(part, "text", ""))

        return " ".join(text_parts)

    @staticmethod
    def extract_text_from_response(response_dict: Mapping[str, object]) -> str:
        """
        Extract text content from A2A response result.

        Args:
            response_dict: A2A response dict with 'result' containing message

        Returns:
            Text from response message parts
        """
        result: Final = response_dict.get("result", {})
        if not isinstance(result, dict):
            return ""

        # Direct message format (A2A spec): detect by explicit kind tag only.
        # The "parts" heuristic is too broad and would match any future result
        # type that happens to include a "parts" field.
        if result.get("kind") == "message":
            return A2ARequestUtils.extract_text_from_message(result)

        message: Final = result.get("message", {})
        return A2ARequestUtils.extract_text_from_message(message)

    @staticmethod
    def get_input_message_from_request(
        request: "SendMessageRequest | SendStreamingMessageRequest",
    ) -> Any:
        """
        Extract the input message from an A2A request.

        Args:
            request: The A2A SendMessageRequest or SendStreamingMessageRequest

        Returns:
            The message object/dict or None
        """
        params: Final = getattr(request, "params", None)
        if params is None:
            return None
        return getattr(params, "message", None)

    @staticmethod
    def count_tokens(text: str) -> int:
        """
        Count tokens in text using litellm.token_counter.

        Args:
            text: Text to count tokens for

        Returns:
            Token count, or 0 if counting fails
        """
        if not text:
            return 0
        try:
            return litellm.token_counter(text=text)
        except Exception:
            verbose_logger.debug("Failed to count tokens")
            return 0

    @staticmethod
    def calculate_usage_from_request_response(
        request: "SendMessageRequest | SendStreamingMessageRequest",
        response_dict: Mapping[str, object],
    ) -> tuple[int, int, int]:
        """
        Calculate token usage from A2A request and response.

        Args:
            request: The A2A SendMessageRequest or SendStreamingMessageRequest
            response_dict: The A2A response as a dict

        Returns:
            Tuple of (prompt_tokens, completion_tokens, total_tokens)
        """
        # Count input tokens. Dump the message to a dict first so extraction hits
        # the dict branch — request-side parts are a2a-sdk Part RootModels whose
        # kind/text live on part.root, which the object branch cannot read. This
        # mirrors how the response side already works (it operates on model_dump).
        input_message = A2ARequestUtils.get_input_message_from_request(request)
        if input_message is not None and hasattr(input_message, "model_dump"):
            input_message = input_message.model_dump(mode="json")
        input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
        prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)

        # Count output tokens
        output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict)
        completion_tokens: Final = A2ARequestUtils.count_tokens(output_text)

        total_tokens: Final = prompt_tokens + completion_tokens

        return prompt_tokens, completion_tokens, total_tokens


def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None:
    message: Final = params.get("message", {})
    if isinstance(message, dict):
        return message.get("contextId")
    return getattr(message, "contextId", None)


def scope_session_to_principal(session_id: str, principal: str | None) -> str:
    """
    Bind a client-supplied A2A contextId to the authenticated principal.

    Without this, two distinct keys authorized for the same agent could set the
    same contextId and read/append to each other's backend memory. The
    principal is hashed (it is already a hashed token) so the raw value is never
    sent to the agent backend, while the original contextId is kept as a suffix
    for operator-side correlation.
    """
    if not principal:
        return session_id
    principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16]
    return f"{principal_prefix}-{session_id}"


# Backwards compatibility aliases
def extract_text_from_a2a_message(message: Any) -> str:
    return A2ARequestUtils.extract_text_from_message(message)


def extract_text_from_a2a_response(response_dict: Mapping[str, object]) -> str:
    return A2ARequestUtils.extract_text_from_response(response_dict)