add exception handling

litellm/a2a_protocol/exception_mapping_utils.py

@@ -0,0 +1,192 @@
+"""
+A2A Protocol Exception Mapping Utils.
+
+Maps A2A SDK exceptions to LiteLLM A2A exception types.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from litellm._logging import verbose_logger
+from litellm.a2a_protocol.card_resolver import (
+    fix_agent_card_url,
+    is_localhost_or_internal_url,
+)
+from litellm.a2a_protocol.exceptions import (
+    A2AAgentCardError,
+    A2AConnectionError,
+    A2AError,
+    A2ALocalhostURLError,
+)
+from litellm.constants import CONNECTION_ERROR_PATTERNS
+
+if TYPE_CHECKING:
+    from a2a.client import A2AClient as A2AClientType
+
+
+# Runtime import
+_A2AClient: Any = None
+try:
+    from a2a.client import A2AClient as _A2AClient
+except ImportError:
+    pass
+
+
+class A2AExceptionCheckers:
+    """
+    Helper class for checking various A2A error conditions.
+    """
+
+    @staticmethod
+    def is_connection_error(error_str: str) -> bool:
+        """
+        Check if an error string indicates a connection error.
+
+        Args:
+            error_str: The error string to check
+
+        Returns:
+            True if the error indicates a connection issue
+        """
+        if not isinstance(error_str, str):
+            return False
+
+        error_str_lower = error_str.lower()
+        return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
+
+    @staticmethod
+    def is_localhost_url(url: Optional[str]) -> bool:
+        """
+        Check if a URL is a localhost/internal URL.
+
+        Args:
+            url: The URL to check
+
+        Returns:
+            True if the URL is localhost/internal
+        """
+        return is_localhost_or_internal_url(url)
+
+    @staticmethod
+    def is_agent_card_error(error_str: str) -> bool:
+        """
+        Check if an error string indicates an agent card error.
+
+        Args:
+            error_str: The error string to check
+
+        Returns:
+            True if the error is related to agent card fetching/parsing
+        """
+        if not isinstance(error_str, str):
+            return False
+
+        error_str_lower = error_str.lower()
+        agent_card_patterns = [
+            "agent card",
+            "agent-card",
+            ".well-known",
+            "card not found",
+            "invalid agent",
+        ]
+        return any(pattern in error_str_lower for pattern in agent_card_patterns)
+
+
+def map_a2a_exception(
+    original_exception: Exception,
+    card_url: Optional[str] = None,
+    api_base: Optional[str] = None,
+    model: Optional[str] = None,
+) -> Exception:
+    """
+    Map an A2A SDK exception to a LiteLLM A2A exception type.
+
+    Args:
+        original_exception: The original exception from the A2A SDK
+        card_url: The URL from the agent card (if available)
+        api_base: The original API base URL
+        model: The model/agent name
+
+    Returns:
+        A mapped LiteLLM A2A exception
+
+    Raises:
+        A2ALocalhostURLError: If the error is a connection error to a localhost URL
+        A2AConnectionError: If the error is a general connection error
+        A2AAgentCardError: If the error is related to agent card issues
+        A2AError: For other A2A-related errors
+    """
+    error_str = str(original_exception)
+
+    # Check for localhost URL connection error (special case - retryable)
+    if (
+        card_url
+        and api_base
+        and A2AExceptionCheckers.is_localhost_url(card_url)
+        and A2AExceptionCheckers.is_connection_error(error_str)
+    ):
+        raise A2ALocalhostURLError(
+            localhost_url=card_url,
+            base_url=api_base,
+            original_error=original_exception,
+            model=model,
+        )
+
+    # Check for agent card errors
+    if A2AExceptionCheckers.is_agent_card_error(error_str):
+        raise A2AAgentCardError(
+            message=error_str,
+            url=api_base,
+            model=model,
+        )
+
+    # Check for general connection errors
+    if A2AExceptionCheckers.is_connection_error(error_str):
+        raise A2AConnectionError(
+            message=error_str,
+            url=card_url or api_base,
+            model=model,
+        )
+
+    # Default: wrap in generic A2AError
+    raise A2AError(
+        message=error_str,
+        model=model,
+    )
+
+
+def handle_a2a_localhost_retry(
+    error: A2ALocalhostURLError,
+    agent_card: Any,
+    a2a_client: "A2AClientType",
+    is_streaming: bool = False,
+) -> "A2AClientType":
+    """
+    Handle A2ALocalhostURLError by fixing the URL and creating a new client.
+
+    This is called when we catch an A2ALocalhostURLError and want to retry
+    with the corrected URL.
+
+    Args:
+        error: The localhost URL error
+        agent_card: The agent card object to fix
+        a2a_client: The current A2A client
+        is_streaming: Whether this is a streaming request (for logging)
+
+    Returns:
+        A new A2A client with the fixed URL
+    """
+    request_type = "streaming " if is_streaming else ""
+    verbose_logger.warning(
+        f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. "
+        f"Agent card contains localhost/internal URL. "
+        f"Retrying with base_url '{error.base_url}'."
+    )
+
+    # Fix the agent card URL
+    fix_agent_card_url(agent_card, error.base_url)
+
+    # Create a new client with the fixed agent card (transport caches URL)
+    return _A2AClient(
+        httpx_client=a2a_client._transport.httpx_client,  # type: ignore[union-attr]
+        agent_card=agent_card,
+    )

litellm/a2a_protocol/exceptions.py

@@ -0,0 +1,150 @@
+"""
+A2A Protocol Exceptions.
+
+Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
+"""
+
+from typing import Optional
+
+import httpx
+
+
+class A2AError(Exception):
+    """
+    Base exception for A2A protocol errors.
+
+    Follows the same pattern as LiteLLM's main exceptions.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        status_code: int = 500,
+        llm_provider: str = "a2a_agent",
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        num_retries: Optional[int] = None,
+    ):
+        self.status_code = status_code
+        self.message = f"litellm.A2AError: {message}"
+        self.llm_provider = llm_provider
+        self.model = model
+        self.litellm_debug_info = litellm_debug_info
+        self.max_retries = max_retries
+        self.num_retries = num_retries
+        self.response = response or httpx.Response(
+            status_code=self.status_code,
+            request=httpx.Request(method="POST", url="https://litellm.ai"),
+        )
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        _message = self.message
+        if self.num_retries:
+            _message += f" LiteLLM Retried: {self.num_retries} times"
+        if self.max_retries:
+            _message += f", LiteLLM Max Retries: {self.max_retries}"
+        return _message
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+class A2AConnectionError(A2AError):
+    """
+    Raised when connection to an A2A agent fails.
+
+    This typically occurs when:
+    - The agent is unreachable
+    - The agent card contains a localhost/internal URL
+    - Network issues prevent connection
+    """
+
+    def __init__(
+        self,
+        message: str,
+        url: Optional[str] = None,
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        num_retries: Optional[int] = None,
+    ):
+        self.url = url
+        super().__init__(
+            message=message,
+            status_code=503,
+            llm_provider="a2a_agent",
+            model=model,
+            response=response,
+            litellm_debug_info=litellm_debug_info,
+            max_retries=max_retries,
+            num_retries=num_retries,
+        )
+
+
+class A2AAgentCardError(A2AError):
+    """
+    Raised when there's an issue with the agent card.
+
+    This includes:
+    - Failed to fetch agent card
+    - Invalid agent card format
+    - Missing required fields
+    """
+
+    def __init__(
+        self,
+        message: str,
+        url: Optional[str] = None,
+        model: Optional[str] = None,
+        response: Optional[httpx.Response] = None,
+        litellm_debug_info: Optional[str] = None,
+    ):
+        self.url = url
+        super().__init__(
+            message=message,
+            status_code=404,
+            llm_provider="a2a_agent",
+            model=model,
+            response=response,
+            litellm_debug_info=litellm_debug_info,
+        )
+
+
+class A2ALocalhostURLError(A2AConnectionError):
+    """
+    Raised when an agent card contains a localhost/internal URL.
+
+    Many A2A agents are deployed with agent cards that contain internal URLs
+    like "http://0.0.0.0:8001/" or "http://localhost:8000/". This error
+    indicates that the URL needs to be corrected and the request should be retried.
+
+    Attributes:
+        localhost_url: The localhost/internal URL found in the agent card
+        base_url: The public base URL that should be used instead
+        original_error: The original connection error that was raised
+    """
+
+    def __init__(
+        self,
+        localhost_url: str,
+        base_url: str,
+        original_error: Optional[Exception] = None,
+        model: Optional[str] = None,
+    ):
+        self.localhost_url = localhost_url
+        self.base_url = base_url
+        self.original_error = original_error
+
+        message = (
+            f"Agent card contains localhost/internal URL '{localhost_url}'. "
+            f"Retrying with base URL '{base_url}'."
+        )
+        super().__init__(
+            message=message,
+            url=localhost_url,
+            model=model,
+        )

litellm/constants.py

@@ -315,6 +315,13 @@ LOCALHOST_URL_PATTERNS: List[str] = [
     "0.0.0.0",
     "[::1]",  # IPv6 localhost
 ]
+# Patterns in error messages that indicate a connection failure
+CONNECTION_ERROR_PATTERNS: List[str] = [
+    "connect",
+    "connection",
+    "network",
+    "refused",
+]
 STREAM_SSE_DONE_STRING: str = "[DONE]"
 STREAM_SSE_DATA_PREFIX: str = "data: "
 ### SPEND TRACKING ###