traia-iatp 0.1.29__py3-none-any.whl

traia_iatp/registry/readmes/REFACTORING_SUMMARY.md

@@ -0,0 +1,191 @@
+# Registry and Search API Refactoring Summary
+
+## Overview
+
+This refactoring separates the concerns between registry management (write operations) and search/query operations (read operations) to eliminate duplication and improve maintainability.
+
+## Changes Made
+
+### 1. `mongodb_registry.py` - Write Operations Only
+
+**Removed Methods:**
+- `query_agents()` - Moved to search API
+- `atlas_search()` - Moved to search API  
+- `vector_search()` - Moved to search API
+- `vector_search_text()` - Moved to search API
+- `query_mcp_servers()` - Moved to search API
+- `atlas_search()` (MCP) - Moved to search API
+- `vector_search()` (MCP) - Moved to search API
+- `vector_search_text()` (MCP) - Moved to search API
+
+**Kept Methods:**
+- `add_utility_agent()` - Registry management
+- `update_health_status()` - Registry management
+- `update_agent_base_url()` - Registry management
+- `add_tags()` - Registry management
+- `remove_agent()` - Registry management
+- `update_utility_agent()` - Registry management
+- `get_statistics()` - Registry statistics (read but specific to registry)
+- `add_mcp_server()` - Registry management
+- `get_mcp_server()` - Simple lookup (kept for backward compatibility)
+
+**Updated Documentation:**
+- Updated module docstring to clarify write-only purpose
+- Updated class docstrings to indicate write operations only
+- Added references to `iatp_search_api.py` for search operations
+
+### 2. `iatp_search_api.py` - Search and Query Operations
+
+**Existing Methods (Unchanged):**
+- `find_utility_agent()` - Search by name, capability, tag, or query
+- `list_utility_agents()` - List agents with filters
+- `search_utility_agents()` - Vector search for agents
+- `find_mcp_server()` - Search MCP servers
+- `list_mcp_servers()` - List MCP servers
+- `search_mcp_servers()` - Vector search for MCP servers
+- `get_mcp_server()` - Get MCP server by name
+
+**Features:**
+- Atlas Search integration
+- Vector search with embeddings
+- Fallback to Atlas Search when vector search fails
+- Read-only operations (no MongoDB write access needed)
+- Cached connections for better performance
+
+### 3. Updated Test Files
+
+**`test_discovery_and_usage.py`:**
+- Updated imports to use `iatp_search_api` instead of `mongodb_registry`
+- Modified `discover_hyperliquid_agent()` to use `search_utility_agents()`
+- Updated `create_tools_from_discovered_agent()` to work with `UtilityAgentInfo` structure
+- Removed registry connection management
+
+**`test_refactoring.py`:**
+- New test script to verify refactoring works correctly
+- Tests registry write operations
+- Tests search API operations
+- Tests discovery test imports
+
+## Benefits
+
+### 1. Clear Separation of Concerns
+- **Registry**: Handles data persistence and lifecycle management
+- **Search API**: Handles discovery and querying
+
+### 2. Reduced Duplication
+- Eliminated duplicate search methods between files
+- Single source of truth for search logic
+
+### 3. Better Performance
+- Search API uses read-only connections
+- Cached connections for repeated queries
+- Optimized for search operations
+
+### 4. Improved Maintainability
+- Easier to modify search logic in one place
+- Clearer responsibilities for each module
+- Better testability
+
+### 5. Enhanced Security
+- Search API can use read-only credentials
+- Registry operations require write permissions
+- Reduced attack surface
+
+## Usage Examples
+
+### Registry Operations (Write)
+```python
+from traia_iatp.registry.mongodb_registry import UtilityAgentRegistry
+
+# Add new agent
+registry = UtilityAgentRegistry()
+await registry.add_utility_agent(agent, tags=["trading", "hyperliquid"])
+
+# Update agent
+await registry.update_health_status(agent_id, is_healthy=True)
+await registry.update_agent_base_url(agent_id, new_url)
+```
+
+### Search Operations (Read)
+```python
+from traia_iatp.registry.iatp_search_api import search_utility_agents, find_utility_agent
+
+# Search for agents
+agents = await search_utility_agents("hyperliquid trading", limit=5)
+
+# Find specific agent
+agent = find_utility_agent(capability="market_data")
+```
+
+### Discovery Test
+```python
+# The discovery test now uses the search API
+from traia_iatp.registry.iatp_search_api import search_utility_agents
+
+# Search for Hyperliquid agent
+search_results = await search_utility_agents(
+    query="hyperliquid trading",
+    limit=5,
+    active_only=True
+)
+```
+
+## Migration Guide
+
+### For Existing Code
+
+1. **If using registry for search/query:**
+   ```python
+   # OLD
+   from traia_iatp.registry.mongodb_registry import UtilityAgentRegistry
+   registry = UtilityAgentRegistry()
+   agents = await registry.query_agents(query="trading")
+   
+   # NEW
+   from traia_iatp.registry.iatp_search_api import search_utility_agents
+   agents = await search_utility_agents("trading")
+   ```
+
+2. **If using registry for write operations:**
+   ```python
+   # No changes needed - these methods are still available
+   from traia_iatp.registry.mongodb_registry import UtilityAgentRegistry
+   registry = UtilityAgentRegistry()
+   await registry.add_utility_agent(agent)
+   ```
+
+### Environment Variables
+
+No changes to environment variables are required. Both modules use the same authentication methods:
+- `MONGODB_CONNECTION_STRING`
+- `MONGODB_USER` + `MONGODB_PASSWORD`
+- `MONGODB_X509_CERT_FILE`
+
+## Testing
+
+Run the refactoring test to verify everything works:
+
+```bash
+cd traia-centralized-backend
+python test_refactoring.py
+```
+
+This will test:
+- Registry write operations
+- Search API operations
+- Discovery test imports
+
+## Future Improvements
+
+1. **Connection Pooling**: Implement connection pooling for better performance
+2. **Caching**: Add Redis caching for frequently accessed data
+3. **Rate Limiting**: Add rate limiting for search operations
+4. **Metrics**: Add metrics collection for search performance
+5. **API Versioning**: Consider API versioning for future changes
+
+## Backward Compatibility
+
+- All existing registry write operations remain unchanged
+- Search operations now use the dedicated search API
+- No breaking changes to existing functionality
+- Clear migration path provided 

traia_iatp/scripts/__init__.py

@@ -0,0 +1,2 @@
+"""IATP utility scripts."""
+

traia_iatp/scripts/create_wallet.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+"""Create an IATPWallet for a user.
+
+This script:
+1. Takes an owner private key as input
+2. Generates a new operator keypair
+3. Calls IATPWalletFactory.createWallet() to deploy an IATPWallet
+4. Returns the operator keys and wallet address
+
+Usage:
+    uv run python -m traia_iatp.scripts.create_wallet \\
+        --owner-key 0x... \\
+        --network sepolia \\
+        [--wait-confirmations 2]
+"""
+
+import argparse
+import json
+import logging
+import sys
+from eth_account import Account
+from web3 import Web3
+
+from traia_iatp.contracts import (
+    get_contract_address,
+    load_contract,
+    ContractName
+)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def create_wallet(
+    owner_private_key: str,
+    network: str = "sepolia",
+    wait_confirmations: int = 2
+) -> dict:
+    """Create an IATP Wallet for an owner.
+    
+    Args:
+        owner_private_key: Owner's private key (0x...)
+        network: Network name (default: "sepolia")
+        wait_confirmations: Number of confirmations to wait (default: 2)
+        
+    Returns:
+        Dict with operator keys and wallet address:
+        {
+            "owner_address": "0x...",
+            "operator_address": "0x...",
+            "operator_private_key": "0x...",
+            "wallet_address": "0x...",
+            "transaction_hash": "0x...",
+            "network": "sepolia"
+        }
+    """
+    logger.info(f"Creating IATP Wallet on {network}")
+    
+    # Create owner account
+    if not owner_private_key.startswith("0x"):
+        owner_private_key = f"0x{owner_private_key}"
+    
+    owner_account = Account.from_key(owner_private_key)
+    logger.info(f"Owner address: {owner_account.address}")
+    
+    # Generate new operator keypair
+    operator_account = Account.create()
+    logger.info(f"Generated operator address: {operator_account.address}")
+    logger.info(f"Operator private key: {operator_account.key.hex()}")
+    
+    # Load IATPWalletFactory contract
+    factory_contract, w3 = load_contract(
+        ContractName.IATP_WALLET_FACTORY,
+        network=network
+    )
+    
+    factory_address = factory_contract.address
+    logger.info(f"IATPWalletFactory address: {factory_address}")
+    
+    # Check owner balance
+    owner_balance = w3.eth.get_balance(owner_account.address)
+    logger.info(f"Owner balance: {w3.from_wei(owner_balance, 'ether')} ETH")
+    
+    if owner_balance == 0:
+        logger.error("Owner has no ETH for gas. Please fund the account.")
+        raise ValueError("Insufficient balance for gas")
+    
+    # Build transaction to create wallet
+    logger.info("Building createWallet transaction...")
+    
+    tx = factory_contract.functions.createWallet(
+        operator_account.address
+    ).build_transaction({
+        'from': owner_account.address,
+        'nonce': w3.eth.get_transaction_count(owner_account.address),
+        'gas': 2000000,  # Estimate gas
+        'gasPrice': w3.eth.gas_price,
+        'chainId': w3.eth.chain_id
+    })
+    
+    # Sign transaction
+    signed_tx = owner_account.sign_transaction(tx)
+    
+    # Send transaction
+    logger.info("Sending transaction...")
+    tx_hash = w3.eth.send_raw_transaction(signed_tx.raw_transaction)
+    logger.info(f"Transaction hash: {tx_hash.hex()}")
+    
+    # Wait for confirmation
+    logger.info(f"Waiting for {wait_confirmations} confirmation(s)...")
+    tx_receipt = w3.eth.wait_for_transaction_receipt(tx_hash, timeout=300)
+    
+    if tx_receipt['status'] != 1:
+        logger.error("Transaction failed!")
+        raise Exception(f"Transaction failed: {tx_hash.hex()}")
+    
+    logger.info(f"Transaction confirmed in block {tx_receipt['blockNumber']}")
+    
+    # Parse WalletCreated event to get wallet address
+    wallet_address = None
+    for log in tx_receipt['logs']:
+        try:
+            event = factory_contract.events.WalletCreated().process_log(log)
+            wallet_address = event['args']['wallet']
+            logger.info(f"✅ IATPWallet created: {wallet_address}")
+            break
+        except:
+            continue
+    
+    if not wallet_address:
+        logger.error("Could not find WalletCreated event in logs")
+        raise Exception("Wallet address not found in transaction logs")
+    
+    # Return result
+    result = {
+        "owner_address": owner_account.address,
+        "operator_address": operator_account.address,
+        "operator_private_key": operator_account.key.hex(),
+        "wallet_address": wallet_address,
+        "transaction_hash": tx_hash.hex(),
+        "block_number": tx_receipt['blockNumber'],
+        "network": network,
+        "chain_id": w3.eth.chain_id
+    }
+    
+    return result
+
+
+def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(
+        description="Create an IATP Wallet via IATPWalletFactory",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Create wallet on Sepolia
+  uv run python -m traia_iatp.scripts.create_wallet \\
+      --owner-key 0x0cdda2d9d744d6d2a3ba4b30ac51227590f0e8f44cf5251bb1ed0941d677c0a4 \\
+      --network sepolia
+  
+  # Create wallet with output to JSON file
+  uv run python -m traia_iatp.scripts.create_wallet \\
+      --owner-key 0x... \\
+      --network sepolia \\
+      --output wallet_config.json
+        """
+    )
+    
+    parser.add_argument(
+        "--owner-key",
+        required=True,
+        help="Owner's private key (0x...)"
+    )
+    
+    parser.add_argument(
+        "--network",
+        default="sepolia",
+        choices=["sepolia", "base-sepolia", "arbitrum-sepolia", "localhost"],
+        help="Network to deploy on (default: sepolia)"
+    )
+    
+    parser.add_argument(
+        "--wait-confirmations",
+        type=int,
+        default=2,
+        help="Number of confirmations to wait (default: 2)"
+    )
+    
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        help="Output JSON file path (prints to stdout if not specified)"
+    )
+    
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging"
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    try:
+        # Create wallet
+        result = create_wallet(
+            owner_private_key=args.owner_key,
+            network=args.network,
+            wait_confirmations=args.wait_confirmations
+        )
+        
+        # Output result
+        output_json = json.dumps(result, indent=2)
+        
+        if args.output:
+            with open(args.output, 'w') as f:
+                f.write(output_json)
+            logger.info(f"✅ Configuration saved to: {args.output}")
+        else:
+            print("\n" + "=" * 80)
+            print("IATP Wallet Created Successfully!")
+            print("=" * 80)
+            print(output_json)
+            print("=" * 80)
+        
+        return 0
+        
+    except Exception as e:
+        logger.error(f"Failed to create wallet: {e}")
+        import traceback
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
+

traia_iatp/server/__init__.py

@@ -0,0 +1,15 @@
+"""IATP server module for creating A2A servers from utility agents."""
+
+from .a2a_server import (
+    UtilityAgencyExecutor,
+    create_a2a_server,
+    create_and_start_a2a_server
+)
+from .iatp_server_agent_generator import IATPServerAgentGenerator
+
+__all__ = [
+    "UtilityAgencyExecutor", 
+    "IATPServerAgentGenerator",
+    "create_a2a_server",
+    "create_and_start_a2a_server",
+]

traia_iatp/server/a2a_server.py

@@ -0,0 +1,219 @@
+"""A2A server implementation for utility agencies using the official a2a-sdk."""
+
+import asyncio
+import logging
+from typing import Dict, Any, Optional, List
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events.event_queue import EventQueue
+from a2a.server.tasks import InMemoryTaskStore
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.types import AgentCard, AgentSkill, AgentCapabilities
+from a2a.utils import new_agent_text_message, new_text_artifact
+from dataclasses import dataclass
+import uvicorn
+
+from ..core.models import UtilityAgent
+from .iatp_server_agent_generator import IATPServerAgentGenerator
+
+logger = logging.getLogger(__name__)
+
+
+class UtilityAgencyExecutor(AgentExecutor):
+    """Agent executor that wraps a utility agency and executes tasks via CrewAI."""
+    
+    def __init__(self, agency: UtilityAgent, agency_generator: IATPServerAgentGenerator):
+        self.agency = agency
+        self.agency_generator = agency_generator
+        self._crew = None
+    
+    async def _get_crew(self):
+        """Lazily build the crew from agency config."""
+        if self._crew is None:
+            self._crew = await self.agency_generator.build_crew_from_agency(self.agency)
+        return self._crew
+    
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Execute a task using the utility agency crew."""
+        try:
+            # Get the user's request from context
+            user_message = context.get_last_user_message()
+            if not user_message:
+                await event_queue.enqueue_event(
+                    new_agent_text_message("No user message provided")
+                )
+                return
+            
+            # Extract text from message parts
+            user_text = ""
+            for part in user_message.parts:
+                if hasattr(part, 'text'):
+                    user_text += part.text
+            
+            if not user_text:
+                event_queue.enqueue_event(
+                    new_agent_text_message("No text content in user message")
+                )
+                return
+            
+            # Build crew and execute
+            crew = await self._get_crew()
+            
+            # Run the crew with the request
+            result = crew.kickoff(inputs={"request": user_text})
+            
+            # Send the result as agent message
+            event_queue.enqueue_event(
+                new_agent_text_message(str(result))
+            )
+            
+        except Exception as e:
+            logger.error(f"Error executing task: {e}")
+            event_queue.enqueue_event(
+                new_agent_text_message(f"Error processing request: {str(e)}")
+            )
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel all tasks."""
+        pass
+
+
+class MCPToolExecutor(AgentExecutor):
+    """Agent executor for individual MCP tool execution."""
+    
+    def __init__(self, agency: UtilityAgent, tool_name: str):
+        self.agency = agency
+        self.tool_name = tool_name
+    
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Execute a specific MCP tool."""
+        try:
+            from ..mcp.client import MCPClient
+            
+            # Get tool arguments from context
+            user_message = context.get_last_user_message()
+            if not user_message:
+                await event_queue.enqueue_event(
+                    new_agent_text_message("No arguments provided")
+                )
+                return
+            
+            # Extract arguments (assuming they're passed as JSON in text)
+            import json
+            arguments = {}
+            for part in user_message.parts:
+                if hasattr(part, 'text'):
+                    try:
+                        arguments = json.loads(part.text)
+                    except:
+                        arguments = {"query": part.text}  # Fallback to simple text
+            
+            # Connect to MCP server and execute tool
+            mcp_client = MCPClient(self.agency.mcp_server)
+            await mcp_client.connect()
+            
+            try:
+                result = await mcp_client.call_tool(self.tool_name, arguments)
+                await event_queue.enqueue_event(
+                    new_agent_text_message(str(result))
+                )
+            finally:
+                await mcp_client.disconnect()
+                
+        except Exception as e:
+            logger.error(f"Error executing MCP tool {self.tool_name}: {e}")
+            await event_queue.enqueue_event(
+                new_agent_text_message(f"Error executing tool: {str(e)}")
+            )
+
+
+def create_a2a_server(
+    agency: UtilityAgent,
+    agency_generator: IATPServerAgentGenerator,
+    host: str = "0.0.0.0",
+    port: int = 8000
+) -> A2AStarletteApplication:
+    """Create an A2A server for a utility agency."""
+    
+    # Create skills from agency capabilities
+    skills = []
+    
+    # Add main processing skill
+    main_skill = AgentSkill(
+        id="process_request",
+        name=f"Process request using {agency.name}",
+        description=f"Process a request using {agency.name} capabilities. {agency.description}",
+        examples=[
+            f"Help me with {cap}" for cap in agency.mcp_server.capabilities[:2]
+        ] if agency.mcp_server.capabilities else ["Process this request for me"]
+    )
+    skills.append(main_skill)
+    
+    # Add individual MCP tool skills
+    for capability in agency.mcp_server.capabilities:
+        skill = AgentSkill(
+            id=f"mcp_{capability}",
+            name=f"Execute {capability}",
+            description=f"Execute {capability} tool on MCP server",
+            examples=[f"Run {capability} with these parameters"]
+        )
+        skills.append(skill)
+    
+    # Create capabilities
+    capabilities = AgentCapabilities(
+        streaming=False,  # Not implementing streaming for now
+        pushNotifications=False,  # Not implementing push notifications
+        stateTransitionHistory=False
+    )
+    
+    # Create agent card
+    agent_card = AgentCard(
+        name=agency.name.replace(" ", "_").lower(),
+        description=agency.description,
+        url=f"http://{host}:{port}",
+        version="1.0.0",
+        capabilities=capabilities,
+        skills=skills,
+        # TODO: Add authentication when AgentAuthentication is available
+        # authentication=AgentAuthentication(schemes=["Bearer"]) if agency.auth_type else None
+    )
+    
+    # Create executor mapping
+    executors = {
+        "process_request": UtilityAgencyExecutor(agency, agency_generator)
+    }
+    
+    # Add MCP tool executors
+    for capability in agency.mcp_server.capabilities:
+        executors[f"mcp_{capability}"] = MCPToolExecutor(agency, capability)
+    
+    # Create task store and request handler
+    task_store = InMemoryTaskStore()
+    request_handler = DefaultRequestHandler(
+        agent_card=agent_card,
+        executors=executors,
+        task_store=task_store
+    )
+    
+    # Create the A2A application
+    app = A2AStarletteApplication(
+        agent_card=agent_card,
+        http_handler=request_handler
+    )
+    
+    return app
+
+
+async def create_and_start_a2a_server(
+    agency: UtilityAgent,
+    agency_generator: IATPServerAgentGenerator,
+    host: str = "0.0.0.0",
+    port: int = 8000
+):
+    """Create and start an A2A server for a utility agency."""
+    app = create_a2a_server(agency, agency_generator, host, port)
+    
+    logger.info(f"Starting A2A server for {agency.name} on {host}:{port}")
+    config = uvicorn.Config(app, host=host, port=port, log_level="info")
+    server = uvicorn.Server(config)
+    await server.serve()