traia-iatp 0.1.29__py3-none-any.whl

traia_iatp/client/d402_a2a_client.py

@@ -0,0 +1,293 @@
+"""D402-enabled A2A client for IATP with payment support."""
+
+import asyncio
+import logging
+import json
+import base64
+from typing import Optional, Dict, Any
+from a2a.client import A2AClient, A2ACardResolver
+from a2a.types import Message, TextPart, TaskState
+import httpx
+
+from ..d402.client import D402IATPClient, create_iatp_payment_client
+from ..d402.models import D402PaymentInfo
+from ..d402.types import d402PaymentRequiredResponse, PaymentRequirements
+
+logger = logging.getLogger(__name__)
+
+
+class D402A2AClient:
+    """A2A client with d402 payment support for IATP.
+    
+    This client automatically handles d402 payments when communicating with
+    utility agents that require payment.
+    """
+    
+    def __init__(
+        self,
+        agent_endpoint: str,
+        payment_client: Optional[D402IATPClient] = None,
+        max_payment_usd: Optional[float] = None
+    ):
+        """Initialize the d402-enabled A2A client.
+        
+        Args:
+            agent_endpoint: URL of the utility agent
+            payment_client: Optional pre-configured D402IATPClient
+            max_payment_usd: Optional maximum payment amount per request
+        """
+        self.agent_endpoint = agent_endpoint
+        self.payment_client = payment_client
+        self.max_payment_usd = max_payment_usd
+        
+        # Resolve agent card
+        self.card_resolver = A2ACardResolver(agent_endpoint)
+        self.agent_card = self.card_resolver.get_agent_card()
+        
+        # Extract d402 payment information if available
+        self.d402_info = self._extract_d402_info()
+        
+        # Initialize A2A client
+        self.a2a_client = A2AClient(
+            agent_card=self.agent_card,
+            credentials=None  # d402 handles payment authorization
+        )
+    
+    def _extract_d402_info(self) -> Optional[D402PaymentInfo]:
+        """Extract d402 payment information from agent card."""
+        metadata = self.agent_card.metadata or {}
+        d402_data = metadata.get("d402")
+        
+        if d402_data and d402_data.get("enabled"):
+            return D402PaymentInfo(**d402_data)
+        return None
+    
+    async def send_message_with_payment(
+        self,
+        message: str,
+        skill_id: Optional[str] = None
+    ) -> str:
+        """Send a message to the agent, automatically handling d402 payment if required.
+        
+        Args:
+            message: The message to send to the agent
+            skill_id: Optional specific skill to invoke
+            
+        Returns:
+            Agent's response text
+            
+        Raises:
+            ValueError: If payment is required but no payment client is configured
+            RuntimeError: If task execution fails
+        """
+        # Prepare the message
+        msg = Message(
+            role="user",
+            parts=[TextPart(text=message)]
+        )
+        
+        # Prepare headers (will add payment header if needed)
+        headers = {}
+        
+        # Try to send the request
+        async with httpx.AsyncClient(timeout=300.0) as http_client:
+            # First attempt without payment to see if it's required
+            try:
+                # Use A2A client's send_task method
+                task = await self.a2a_client.send_task(
+                    id=str(asyncio.get_event_loop().time()),
+                    message=msg
+                )
+                
+                # Extract response
+                return self._extract_response(task)
+                
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 402:
+                    # Payment required - handle d402 flow
+                    return await self._handle_payment_required(
+                        http_client, e.response, message, skill_id
+                    )
+                else:
+                    raise
+    
+    async def _handle_payment_required(
+        self,
+        http_client: httpx.AsyncClient,
+        response: httpx.Response,
+        message: str,
+        skill_id: Optional[str]
+    ) -> str:
+        """Handle 402 Payment Required response.
+        
+        Args:
+            http_client: HTTP client for making requests
+            response: 402 response with payment requirements
+            message: Original message to send
+            skill_id: Optional skill ID
+            
+        Returns:
+            Agent's response after successful payment
+            
+        Raises:
+            ValueError: If payment client not configured or requirements not met
+        """
+        if not self.payment_client:
+            raise ValueError(
+                "Payment is required but no payment client configured. "
+                "Please provide a payment_client when initializing D402A2AClient."
+            )
+        
+        # Parse payment requirements
+        try:
+            payment_response = d402PaymentRequiredResponse(**response.json())
+        except Exception as e:
+            raise ValueError(f"Invalid payment requirements: {e}")
+        
+        # Select payment requirements
+        try:
+            selected_requirements = self.payment_client.select_payment_requirements(
+                accepts=payment_response.accepts,
+                network_filter=None,  # Accept any network
+                scheme_filter="exact"  # Only support exact scheme
+            )
+        except Exception as e:
+            raise ValueError(f"Cannot satisfy payment requirements: {e}")
+        
+        # Check if amount is within maximum
+        if self.max_payment_usd:
+            # Convert to USD (assuming USDC with 6 decimals)
+            amount_usd = int(selected_requirements.max_amount_required) / 1_000_000
+            if amount_usd > self.max_payment_usd:
+                raise ValueError(
+                    f"Payment amount ${amount_usd:.2f} exceeds maximum ${self.max_payment_usd:.2f}"
+                )
+        
+        # Create payment header
+        payment_header = self.payment_client.create_payment_header(
+            payment_requirements=selected_requirements,
+            d402_version=1
+        )
+        
+        # Retry request with payment header
+        headers = {"X-PAYMENT": payment_header}
+        
+        # Make the A2A request with payment
+        msg = Message(
+            role="user",
+            parts=[TextPart(text=message)]
+        )
+        
+        # Send task with payment (we need to make raw HTTP request with headers)
+        # Since A2AClient doesn't support custom headers, we'll make the request directly
+        request_data = {
+            "jsonrpc": "2.0",
+            "id": str(asyncio.get_event_loop().time()),
+            "method": "message/send",
+            "params": {
+                "message": msg.model_dump()
+            }
+        }
+        
+        response = await http_client.post(
+            self.agent_endpoint,
+            json=request_data,
+            headers={"Content-Type": "application/json", **headers}
+        )
+        
+        if response.status_code == 402:
+            raise RuntimeError(f"Payment failed: {response.json().get('error', 'Unknown error')}")
+        
+        response.raise_for_status()
+        
+        # Parse JSON-RPC response
+        result = response.json()
+        if "error" in result:
+            raise RuntimeError(f"Agent error: {result['error']}")
+        
+        # Extract the response text
+        task_result = result.get("result", {})
+        if isinstance(task_result, dict):
+            messages = task_result.get("messages", [])
+            for msg in messages:
+                if msg.get("role") == "agent":
+                    parts = msg.get("parts", [])
+                    return " ".join(
+                        part.get("text", "")
+                        for part in parts
+                        if part.get("type") == "text"
+                    )
+        
+        return str(task_result)
+    
+    def _extract_response(self, task) -> str:
+        """Extract text response from task result."""
+        if task.status.state == TaskState.COMPLETED:
+            # Look for agent's response in messages
+            for msg in task.messages:
+                if msg.role == "agent":
+                    response_text = ""
+                    for part in msg.parts:
+                        if hasattr(part, 'text'):
+                            response_text += part.text
+                    return response_text
+            
+            # If no agent message, check artifacts
+            if task.artifacts:
+                response_text = ""
+                for artifact in task.artifacts:
+                    for part in artifact.parts:
+                        if hasattr(part, 'text'):
+                            response_text += part.text
+                return response_text
+            
+            return "Task completed but no response found"
+        else:
+            raise RuntimeError(f"Task failed with state: {task.status.state}")
+
+
+def create_d402_a2a_client(
+    agent_endpoint: str,
+    payment_private_key: Optional[str] = None,
+    max_payment_usd: Optional[float] = 10.0,
+    agent_contract_address: Optional[str] = None
+) -> D402A2AClient:
+    """Convenience function to create an d402-enabled A2A client.
+    
+    Args:
+        agent_endpoint: URL of the utility agent
+        payment_private_key: Optional private key for payments (hex encoded)
+        max_payment_usd: Maximum payment per request in USD
+        agent_contract_address: Optional client agent contract address
+        
+    Returns:
+        Configured D402A2AClient
+        
+    Example:
+        # Create client with payment support
+        client = create_d402_a2a_client(
+            agent_endpoint="https://agent.example.com",
+            payment_private_key="0x...",
+            max_payment_usd=5.0
+        )
+        
+        # Send message (automatically handles payment if required)
+        response = await client.send_message_with_payment(
+            "Analyze sentiment of: 'Stock prices are rising'"
+        )
+        print(response)
+    """
+    payment_client = None
+    if payment_private_key:
+        payment_client = create_iatp_payment_client(
+            private_key=payment_private_key,
+            max_value_usd=max_payment_usd,
+            agent_contract_address=agent_contract_address
+        )
+    
+    return D402A2AClient(
+        agent_endpoint=agent_endpoint,
+        payment_client=payment_client,
+        max_payment_usd=max_payment_usd
+    )
+

traia_iatp/client/grpc_a2a_tools.py

@@ -0,0 +1,349 @@
+"""gRPC-based A2A client for high-performance agent communication."""
+
+import asyncio
+import logging
+from typing import Dict, Any, Optional, List, Union, AsyncIterator
+from datetime import datetime
+import json
+import uuid
+from contextlib import asynccontextmanager
+
+import grpc
+from grpc import aio
+from crewai.tools import BaseTool
+from pydantic import Field, BaseModel
+
+# Note: These imports assume the A2A proto files have been compiled
+# In practice, you'd need to generate these from the A2A .proto definitions
+# from a2a.proto import a2a_pb2, a2a_pb2_grpc
+
+logger = logging.getLogger(__name__)
+
+
+class GrpcA2AConfig(BaseModel):
+    """Configuration for gRPC A2A tools."""
+    endpoint: str = Field(description="The gRPC endpoint (host:port)")
+    agency_name: str = Field(description="Name of the utility agency")
+    agency_description: str = Field(description="Description of what the agency does")
+    timeout: int = Field(default=300, description="Timeout in seconds for gRPC calls")
+    retry_attempts: int = Field(default=3, description="Number of retry attempts")
+    
+    # gRPC specific settings
+    max_message_length: int = Field(default=4 * 1024 * 1024, description="Max message size (4MB default)")
+    keepalive_time_ms: int = Field(default=10000, description="Keepalive ping interval")
+    keepalive_timeout_ms: int = Field(default=5000, description="Keepalive timeout")
+    max_concurrent_streams: int = Field(default=100, description="Max concurrent gRPC streams")
+    
+    # Connection pool settings
+    pool_size: int = Field(default=10, description="Number of gRPC channels in pool")
+    use_tls: bool = Field(default=False, description="Whether to use TLS")
+    tls_cert_path: Optional[str] = Field(default=None, description="Path to TLS certificate")
+
+
+class GrpcChannelPool:
+    """Manages a pool of gRPC channels for load distribution."""
+    
+    def __init__(self, config: GrpcA2AConfig):
+        self.config = config
+        self.channels: List[aio.Channel] = []
+        self.current_index = 0
+        self._lock = asyncio.Lock()
+    
+    async def initialize(self):
+        """Initialize the channel pool."""
+        options = [
+            ('grpc.max_send_message_length', self.config.max_message_length),
+            ('grpc.max_receive_message_length', self.config.max_message_length),
+            ('grpc.keepalive_time_ms', self.config.keepalive_time_ms),
+            ('grpc.keepalive_timeout_ms', self.config.keepalive_timeout_ms),
+            ('grpc.http2.max_concurrent_streams', self.config.max_concurrent_streams),
+            ('grpc.enable_http_proxy', 0),  # Disable proxy for direct connection
+        ]
+        
+        for _ in range(self.config.pool_size):
+            if self.config.use_tls:
+                # Load TLS credentials
+                if self.config.tls_cert_path:
+                    with open(self.config.tls_cert_path, 'rb') as f:
+                        credentials = grpc.ssl_channel_credentials(f.read())
+                else:
+                    credentials = grpc.ssl_channel_credentials()
+                
+                channel = aio.secure_channel(
+                    self.config.endpoint,
+                    credentials,
+                    options=options
+                )
+            else:
+                channel = aio.insecure_channel(
+                    self.config.endpoint,
+                    options=options
+                )
+            
+            self.channels.append(channel)
+        
+        logger.info(f"Initialized gRPC channel pool with {self.config.pool_size} channels")
+    
+    async def get_channel(self) -> aio.Channel:
+        """Get the next available channel (round-robin)."""
+        async with self._lock:
+            channel = self.channels[self.current_index]
+            self.current_index = (self.current_index + 1) % len(self.channels)
+            return channel
+    
+    async def close(self):
+        """Close all channels in the pool."""
+        for channel in self.channels:
+            await channel.close()
+        self.channels.clear()
+
+
+class GrpcA2ATool(BaseTool):
+    """gRPC-based A2A tool for high-performance agent communication."""
+    
+    name: str
+    description: str
+    config: GrpcA2AConfig
+    _channel_pool: Optional[GrpcChannelPool] = None
+    _initialized: bool = False
+    
+    def __init__(self, config: GrpcA2AConfig):
+        """Initialize the gRPC A2A tool."""
+        tool_name = f"grpc_a2a_{config.agency_name.replace(' ', '_').replace('-', '_').lower()}"
+        
+        super().__init__(
+            name=tool_name,
+            description=f"Use {config.agency_name} via gRPC A2A: {config.agency_description}",
+            config=config
+        )
+        self._channel_pool = None
+        self._initialized = False
+    
+    async def _ensure_initialized(self):
+        """Ensure the gRPC channel pool is initialized."""
+        if not self._initialized:
+            self._channel_pool = GrpcChannelPool(self.config)
+            await self._channel_pool.initialize()
+            self._initialized = True
+    
+    async def _execute_unary(self, request: str, context: Optional[Dict[str, Any]] = None) -> str:
+        """Execute a unary (request-response) gRPC call."""
+        await self._ensure_initialized()
+        
+        # Get a channel from the pool
+        channel = await self._channel_pool.get_channel()
+        
+        # Create stub (lightweight, can be created per request)
+        # stub = a2a_pb2_grpc.A2AServiceStub(channel)
+        
+        # For now, return a placeholder since we don't have the actual proto files
+        # In a real implementation, you would:
+        # 1. Create the request message
+        # 2. Call the appropriate gRPC method
+        # 3. Handle the response
+        
+        return f"gRPC call to {self.config.agency_name} with request: {request}"
+    
+    async def _execute_streaming(self, request: str, context: Optional[Dict[str, Any]] = None) -> AsyncIterator[str]:
+        """Execute a server-streaming gRPC call."""
+        await self._ensure_initialized()
+        
+        # Get a channel from the pool
+        channel = await self._channel_pool.get_channel()
+        
+        # Create stub
+        # stub = a2a_pb2_grpc.A2AServiceStub(channel)
+        
+        # For demonstration, yield simulated chunks
+        for i in range(5):
+            yield f"Chunk {i} from {self.config.agency_name}"
+            await asyncio.sleep(0.1)
+    
+    async def _execute_with_retry(self, request: str, context: Optional[Dict[str, Any]] = None) -> str:
+        """Execute request with retry logic."""
+        last_error = None
+        
+        for attempt in range(self.config.retry_attempts):
+            try:
+                return await self._execute_unary(request, context)
+                
+            except grpc.RpcError as e:
+                last_error = f"gRPC error: {e.code()}: {e.details()}"
+                logger.warning(f"Attempt {attempt + 1} failed: {last_error}")
+                
+                # Check if retryable
+                if e.code() in [
+                    grpc.StatusCode.UNAVAILABLE,
+                    grpc.StatusCode.DEADLINE_EXCEEDED,
+                    grpc.StatusCode.INTERNAL
+                ]:
+                    # Wait before retry with exponential backoff
+                    if attempt < self.config.retry_attempts - 1:
+                        await asyncio.sleep(2 ** attempt)
+                else:
+                    # Non-retryable error
+                    break
+                    
+            except Exception as e:
+                last_error = str(e)
+                logger.warning(f"Attempt {attempt + 1} failed: {e}")
+                
+                if attempt < self.config.retry_attempts - 1:
+                    await asyncio.sleep(2 ** attempt)
+        
+        return f"Failed after {self.config.retry_attempts} attempts. Last error: {last_error}"
+    
+    async def _arun(self, request: str, **kwargs) -> str:
+        """Async execution of the tool."""
+        try:
+            context = kwargs.get('context', {})
+            return await self._execute_with_retry(request, context)
+        except Exception as e:
+            logger.error(f"Error in gRPC A2A tool {self.name}: {e}")
+            return f"Error: {str(e)}"
+    
+    def _run(self, request: str, **kwargs) -> str:
+        """Sync execution of the tool."""
+        try:
+            loop = asyncio.get_running_loop()
+            return asyncio.run_coroutine_threadsafe(
+                self._arun(request, **kwargs),
+                loop
+            ).result()
+        except RuntimeError:
+            return asyncio.run(self._arun(request, **kwargs))
+    
+    async def stream(self, request: str, **kwargs) -> AsyncIterator[str]:
+        """Stream responses using server-streaming gRPC."""
+        try:
+            context = kwargs.get('context', {})
+            async for chunk in self._execute_streaming(request, context):
+                yield chunk
+        except Exception as e:
+            logger.error(f"Error in gRPC streaming: {e}")
+            yield f"Streaming error: {str(e)}"
+    
+    async def close(self):
+        """Close the gRPC channel pool."""
+        if self._channel_pool:
+            await self._channel_pool.close()
+            self._channel_pool = None
+            self._initialized = False
+        logger.info(f"Closed gRPC connections for {self.name}")
+    
+    def __del__(self):
+        """Cleanup gRPC channels when tool is destroyed."""
+        if self._channel_pool:
+            try:
+                loop = asyncio.get_running_loop()
+                asyncio.create_task(self.close())
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                loop.run_until_complete(self.close())
+                loop.close()
+            except Exception:
+                pass
+
+
+class GrpcA2AToolkit:
+    """Toolkit for creating gRPC-based A2A tools."""
+    
+    @staticmethod
+    def create_tool_from_endpoint(
+        endpoint: str,
+        name: str,
+        description: str,
+        timeout: int = 300,
+        retry_attempts: int = 3,
+        pool_size: int = 10,
+        use_tls: bool = False,
+        tls_cert_path: Optional[str] = None
+    ) -> GrpcA2ATool:
+        """Create a gRPC A2A tool from an endpoint."""
+        config = GrpcA2AConfig(
+            endpoint=endpoint,
+            agency_name=name,
+            agency_description=description,
+            timeout=timeout,
+            retry_attempts=retry_attempts,
+            pool_size=pool_size,
+            use_tls=use_tls,
+            tls_cert_path=tls_cert_path
+        )
+        
+        return GrpcA2ATool(config)
+    
+    @staticmethod
+    def create_tools_from_endpoints(
+        endpoints: List[Dict[str, Any]],
+        default_timeout: int = 300,
+        default_retry_attempts: int = 3,
+        default_pool_size: int = 10
+    ) -> List[GrpcA2ATool]:
+        """Create multiple gRPC A2A tools."""
+        tools = []
+        
+        for ep_config in endpoints:
+            tool = GrpcA2AToolkit.create_tool_from_endpoint(
+                endpoint=ep_config["endpoint"],
+                name=ep_config["name"],
+                description=ep_config["description"],
+                timeout=ep_config.get("timeout", default_timeout),
+                retry_attempts=ep_config.get("retry_attempts", default_retry_attempts),
+                pool_size=ep_config.get("pool_size", default_pool_size),
+                use_tls=ep_config.get("use_tls", False),
+                tls_cert_path=ep_config.get("tls_cert_path")
+            )
+            tools.append(tool)
+            logger.info(f"Created gRPC A2A tool: {tool.name}")
+        
+        return tools
+
+
+# Example usage with both HTTP/2 and gRPC
+async def compare_protocols():
+    """Compare HTTP/2 and gRPC performance."""
+    from .crewai_a2a_tools import A2AToolkit as HttpA2AToolkit
+    
+    # Create HTTP/2 tool
+    http_tool = HttpA2AToolkit.create_tool_from_endpoint(
+        endpoint="http://localhost:8000",
+        name="Trading Agency HTTP",
+        description="Trading via HTTP/2"
+    )
+    
+    # Create gRPC tool
+    grpc_tool = GrpcA2AToolkit.create_tool_from_endpoint(
+        endpoint="localhost:50051",
+        name="Trading Agency gRPC",
+        description="Trading via gRPC",
+        pool_size=20  # More channels for high load
+    )
+    
+    # Run parallel requests
+    import time
+    
+    # HTTP/2 test
+    start = time.time()
+    http_tasks = [
+        http_tool._arun(f"Get price for BTC {i}")
+        for i in range(100)
+    ]
+    await asyncio.gather(*http_tasks)
+    http_time = time.time() - start
+    
+    # gRPC test
+    start = time.time()
+    grpc_tasks = [
+        grpc_tool._arun(f"Get price for BTC {i}")
+        for i in range(100)
+    ]
+    await asyncio.gather(*grpc_tasks)
+    grpc_time = time.time() - start
+    
+    print(f"HTTP/2 time: {http_time:.2f}s")
+    print(f"gRPC time: {grpc_time:.2f}s")
+    
+    # Cleanup
+    await http_tool.close()
+    await grpc_tool.close() 

traia_iatp/client/root_path_a2a_client.py

@@ -0,0 +1 @@
+ 

traia_iatp/contracts/__init__.py

@@ -0,0 +1,12 @@
+"""IATP Contracts utilities."""
+
+from .wallet_creator import create_wallet, get_contract_config
+from .iatp_contracts_config import get_contract_address, get_contract_abi, get_rpc_url
+
+__all__ = [
+    "create_wallet",
+    "get_contract_config",
+    "get_contract_address",
+    "get_contract_abi",
+    "get_rpc_url"
+]