chore: initial public snapshot for github upload

llm-gateway-competitors/litellm-wheel-src/litellm/a2a_protocol/client.py

@@ -0,0 +1,109 @@
+"""
+LiteLLM A2A Client class.
+
+Provides a class-based interface for A2A agent invocation.
+"""
+
+from typing import TYPE_CHECKING, AsyncIterator, Dict, Optional
+
+from litellm.types.agents import LiteLLMSendMessageResponse
+
+if TYPE_CHECKING:
+    from a2a.client import A2AClient as A2AClientType
+    from a2a.types import (
+        AgentCard,
+        SendMessageRequest,
+        SendStreamingMessageRequest,
+        SendStreamingMessageResponse,
+    )
+
+
+class A2AClient:
+    """
+    LiteLLM wrapper for A2A agent invocation.
+
+    Creates the underlying A2A client once on first use and reuses it.
+
+    Example:
+        ```python
+        from litellm.a2a_protocol import A2AClient
+        from a2a.types import SendMessageRequest, MessageSendParams
+        from uuid import uuid4
+
+        client = A2AClient(base_url="http://localhost:10001")
+
+        request = SendMessageRequest(
+            id=str(uuid4()),
+            params=MessageSendParams(
+                message={
+                    "role": "user",
+                    "parts": [{"kind": "text", "text": "Hello!"}],
+                    "messageId": uuid4().hex,
+                }
+            )
+        )
+        response = await client.send_message(request)
+        ```
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 60.0,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize the A2A client wrapper.
+
+        Args:
+            base_url: The base URL of the A2A agent (e.g., "http://localhost:10001")
+            timeout: Request timeout in seconds (default: 60.0)
+            extra_headers: Optional additional headers to include in requests
+        """
+        self.base_url = base_url
+        self.timeout = timeout
+        self.extra_headers = extra_headers
+        self._a2a_client: Optional["A2AClientType"] = None
+
+    async def _get_client(self) -> "A2AClientType":
+        """Get or create the underlying A2A client."""
+        if self._a2a_client is None:
+            from litellm.a2a_protocol.main import create_a2a_client
+
+            self._a2a_client = await create_a2a_client(
+                base_url=self.base_url,
+                timeout=self.timeout,
+                extra_headers=self.extra_headers,
+            )
+        return self._a2a_client
+
+    async def get_agent_card(self) -> "AgentCard":
+        """Fetch the agent card from the server."""
+        from litellm.a2a_protocol.main import aget_agent_card
+
+        return await aget_agent_card(
+            base_url=self.base_url,
+            timeout=self.timeout,
+            extra_headers=self.extra_headers,
+        )
+
+    async def send_message(
+        self, request: "SendMessageRequest"
+    ) -> LiteLLMSendMessageResponse:
+        """Send a message to the A2A agent."""
+        from litellm.a2a_protocol.main import asend_message
+
+        a2a_client = await self._get_client()
+        return await asend_message(a2a_client=a2a_client, request=request)
+
+    async def send_message_streaming(
+        self, request: "SendStreamingMessageRequest"
+    ) -> AsyncIterator["SendStreamingMessageResponse"]:
+        """Send a streaming message to the A2A agent."""
+        from litellm.a2a_protocol.main import asend_message_streaming
+
+        a2a_client = await self._get_client()
+        async for chunk in asend_message_streaming(
+            a2a_client=a2a_client, request=request
+        ):
+            yield chunk