traia-iatp 0.1.1__py3-none-any.whl

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/core/__init__.py

@@ -0,0 +1,43 @@
+"""Core data models for IATP."""
+
+from .models import (
+    # Core server models
+    MCPServer,
+    MCPServerType,
+    
+    # Agent models  
+    AgentSkill,
+    AgentCapabilities,
+    AgentCard,
+    
+    # IATP protocol models
+    IATPEndpoints,
+    IATPRequest,
+    IATPResponse,
+    
+    # Utility agent models
+    UtilityAgent,
+    UtilityAgentStatus,
+    UtilityAgentRegistryEntry,
+)
+
+__all__ = [
+    # Core server models
+    "MCPServer",
+    "MCPServerType",
+    
+    # Agent models
+    "AgentSkill", 
+    "AgentCapabilities",
+    "AgentCard",
+    
+    # IATP protocol models
+    "IATPEndpoints",
+    "IATPRequest",
+    "IATPResponse",
+    
+    # Utility agent models
+    "UtilityAgent",
+    "UtilityAgentStatus", 
+    "UtilityAgentRegistryEntry",
+]

traia_iatp/core/models.py

@@ -0,0 +1,161 @@
+"""Core data models for IATP."""
+
+from datetime import datetime
+from typing import Optional, Dict, Any, List
+from pydantic import BaseModel, Field, HttpUrl
+from enum import Enum
+
+
+class MCPServerType(str, Enum):
+    """Types of MCP servers."""
+    STREAMABLE_HTTP = "streamable-http"
+
+
+class MCPServer(BaseModel):
+    """MCP Server specification."""
+    id: Optional[str] = Field(default=None, description="Unique identifier")
+    name: str = Field(..., description="Name of the MCP server")
+    url: str = Field(..., description="URL or path to connect to the MCP server")
+    server_type: MCPServerType = Field(default=MCPServerType.STREAMABLE_HTTP, description="Type of MCP server connection")
+    description: str = Field(..., description="Description of what the server enables")
+    capabilities: List[str] = Field(default_factory=list, description="List of capabilities/APIs")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+    
+    class Config:
+        json_encoders = {HttpUrl: str}
+
+
+class AgentSkill(BaseModel):
+    """IATP Agent skill definition."""
+    id: str = Field(..., description="Unique skill identifier")
+    name: str = Field(..., description="Human-readable skill name")
+    description: str = Field(..., description="Detailed skill description")
+    examples: List[str] = Field(default_factory=list, description="Example usage patterns")
+    input_modes: List[str] = Field(default_factory=list, description="Supported input modes")
+    output_modes: List[str] = Field(default_factory=list, description="Supported output modes")
+    tags: List[str] = Field(default_factory=list, description="Tags for categorization")
+
+
+class AgentCapabilities(BaseModel):
+    """IATP Agent capabilities."""
+    streaming: bool = Field(default=False, description="Supports SSE streaming")
+    push_notifications: bool = Field(default=False, description="Supports push notifications")
+    state_transition_history: bool = Field(default=False, description="Maintains state history")
+    custom_features: Dict[str, Any] = Field(default_factory=dict, description="Custom capabilities")
+
+
+class AgentCard(BaseModel):
+    """IATP Agent card for discovery and initialization."""
+    name: str = Field(..., description="Agent name")
+    description: str = Field(..., description="Agent description")
+    version: str = Field(..., description="Agent version")
+    skills: List[AgentSkill] = Field(default_factory=list, description="Available skills")
+    capabilities: AgentCapabilities = Field(default_factory=AgentCapabilities, description="Agent capabilities")
+    default_input_modes: List[str] = Field(default_factory=list, description="Default input modes")
+    default_output_modes: List[str] = Field(default_factory=list, description="Default output modes")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+
+
+class IATPEndpoints(BaseModel):
+    """IATP server endpoints configuration.
+    
+    Note: The A2A protocol only defines a minimal set of endpoints:
+    - Root path (/) for JSON-RPC
+    - /.well-known/agent.json for agent card
+    - /a2a/tasks/* for SSE subscriptions (if streaming is supported)
+    
+    Health and info endpoints are NOT part of the A2A protocol standard.
+    """
+    base_url: str = Field(..., description="Base URL of the IATP server")
+    iatp_endpoint: str = Field(..., description="Main IATP JSON-RPC endpoint (usually at root path)")
+    streaming_endpoint: Optional[str] = Field(None, description="SSE streaming endpoint (same as iatp_endpoint when supported)")
+    health_endpoint: Optional[str] = Field(None, description="Health check endpoint (not part of A2A protocol)")
+    info_endpoint: Optional[str] = Field(None, description="Agent info endpoint (not part of A2A protocol)")
+    agent_card_endpoint: str = Field(..., description="Agent card endpoint (.well-known/agent.json)")
+    subscribe_endpoint: Optional[str] = Field(None, description="SSE subscription endpoint (/a2a/tasks/subscribe)")
+    resubscribe_endpoint: Optional[str] = Field(None, description="SSE resubscription endpoint (/a2a/tasks/resubscribe)")
+
+
+class UtilityAgentStatus(str, Enum):
+    """Status of a utility agent."""
+    GENERATED = "generated"
+    DEPLOYING = "deploying"
+    DEPLOYED = "deployed"
+    RUNNING = "running"
+    STOPPED = "stopped"
+    ERROR = "error"
+
+
+#this is the model for internal management of utility agents by their creator and traia protocol (should be persisted into db)
+class UtilityAgent(BaseModel):
+    """Utility Agent configuration and metadata."""
+    id: str = Field(..., description="Unique identifier for the agent")
+    name: str = Field(..., description="Name of the utility agent")
+    description: str = Field(..., description="Description of the agent's purpose")
+    mcp_server_id: str = Field(..., description="ID of the associated MCP server")
+    
+    # IATP specific fields
+    agent_card: Optional[AgentCard] = Field(None, description="IATP agent card for discovery")
+    endpoints: Optional[IATPEndpoints] = Field(None, description="IATP endpoints configuration")
+    
+    capabilities: List[str] = Field(default_factory=list, description="List of exposed capabilities")
+    status: UtilityAgentStatus = Field(default=UtilityAgentStatus.GENERATED)
+    code_path: Optional[str] = Field(None, description="Path to generated code")
+    docker_image: Optional[str] = Field(None, description="Docker image name when built")
+    github_repo: Optional[str] = Field(None, description="GitHub repository URL")
+    cloud_run_url: Optional[str] = Field(None, description="Cloud Run deployment URL")
+    
+    # Search and discovery
+    search_text: Optional[str] = Field(None, description="Concatenated searchable text")
+    tags: List[str] = Field(default_factory=list, description="Tags for search and categorization")
+    
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    updated_at: datetime = Field(default_factory=datetime.utcnow)
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+    
+    class Config:
+        json_encoders = {HttpUrl: str, datetime: lambda v: v.isoformat()}
+
+
+#this is a model for registry discovery that will be used by the mongodb indexes
+class UtilityAgentRegistryEntry(BaseModel):
+    """Registry entry for a deployed utility agent."""
+    agent_id: str = Field(..., description="ID of the utility agent")
+    name: str = Field(..., description="Name for discovery")
+    description: str = Field(..., description="Description for search")
+    
+    # Base URL for the agent - all endpoints are derived from this
+    base_url: Optional[str] = Field(None, description="Base URL of the deployed agent")
+    
+    # Enhanced IATP discovery fields
+    agent_card: Optional[AgentCard] = Field(None, description="IATP agent card")
+    endpoints: Optional[IATPEndpoints] = Field(None, description="IATP endpoints")
+    
+    capabilities: List[str] = Field(..., description="List of capabilities")
+    skills: List[AgentSkill] = Field(default_factory=list, description="Detailed skills from agent card")
+    tags: List[str] = Field(default_factory=list, description="Tags for search")
+    
+    # Search optimization
+    search_text: Optional[str] = Field(None, description="Full text for search")
+    
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+    registered_at: datetime = Field(default_factory=datetime.utcnow)
+    last_health_check: Optional[datetime] = Field(None, description="Last successful health check")
+    is_active: bool = Field(default=True, description="Whether the agent is active")
+    
+    class Config:
+        json_encoders = {HttpUrl: str, datetime: lambda v: v.isoformat()}
+
+
+class IATPRequest(BaseModel):
+    """IATP protocol request."""
+    action: str = Field(..., description="Action to perform")
+    parameters: Dict[str, Any] = Field(default_factory=dict, description="Action parameters")
+    context: Dict[str, Any] = Field(default_factory=dict, description="Request context")
+
+
+class IATPResponse(BaseModel):
+    """IATP protocol response."""
+    result: Any = Field(..., description="Result of the action")
+    status: str = Field(..., description="Status of the response")
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Response metadata") 

traia_iatp/mcp/__init__.py

@@ -0,0 +1,15 @@
+"""MCP (Model Context Protocol) integration module."""
+
+from .client import MCPClient
+from .mcp_agent_template import MCPServerConfig, MCPAgentBuilder, run_with_mcp_tools, MCPServerInfo
+from .traia_mcp_adapter import TraiaMCPAdapter, create_mcp_adapter
+
+__all__ = [
+    "MCPClient",
+    "MCPServerConfig",
+    "MCPAgentBuilder",
+    "run_with_mcp_tools",
+    "MCPServerInfo",
+    "TraiaMCPAdapter",
+    "create_mcp_adapter",
+]