traia-iatp 0.1.2__py3-none-any.whl → 0.1.67__py3-none-any.whl

traia_iatp/mcp/traia_mcp_adapter.py

@@ -3,10 +3,12 @@
 Traia MCP Adapter
 
 A custom MCP adapter that extends CrewAI's MCPServerAdapter to support
-passing headers (like Authorization) to streamable-http MCP servers.
+passing headers (like Authorization) and x402 payment protocol to streamable-http MCP servers.
 
-This adapter transparently handles both authenticated and non-authenticated
-MCP connections by injecting headers when provided.
+This adapter transparently handles:
+- Authenticated connections (via headers)
+- Payment-based connections (via x402/d402 protocol)
+- Standard connections (no auth or payment)
 """
 
 import logging
@@ -19,6 +21,15 @@ from crewai.tools import BaseTool
 
 logger = logging.getLogger(__name__)
 
+# Try to import d402 payment hooks from IATP
+try:
+    from eth_account import Account
+    from traia_iatp.d402.clients.httpx import d402_payment_hooks
+    D402_AVAILABLE = True
+except ImportError:
+    D402_AVAILABLE = False
+    logger.warning("d402 payment hooks not available. Ensure traia_iatp.d402 is installed.")
+
 
 class TraiaMCPAdapter(MCPServerAdapter):
     """
@@ -55,31 +66,53 @@ class TraiaMCPAdapter(MCPServerAdapter):
         *tool_names: str
     ):
         """
-        Initialize the adapter with optional header support.
+        Initialize the adapter with optional header and x402 payment support.
         
         Args:
             server_params: Server configuration. For streamable-http, can include:
                           - url: The MCP server URL
                           - transport: "streamable-http"
                           - headers: Optional dict of headers to include in requests
+                          - x402_account: Optional eth_account.Account for x402 payments
+                          - x402_max_value: Optional max payment value
             *tool_names: Optional tool names to filter
         """
         # Handle dict params
         if isinstance(server_params, dict):
-            # Extract headers if present (don't modify original)
+            # Extract headers and d402 config if present (don't modify original)
             server_params_copy = server_params.copy()
             self._auth_headers = server_params_copy.pop("headers", {})
+            self._d402_account = server_params_copy.pop("d402_account", None)
+            self._d402_max_value = server_params_copy.pop("d402_max_value", None)
+            
+            # IMPORTANT: Also remove any x402_* parameters that shouldn't go to MCP library
+            # The MCP library's streamablehttp_client() doesn't accept these
+            server_params_copy.pop("x402_account", None)
+            server_params_copy.pop("x402_max_value", None)
             
-            if self._auth_headers:
+            # Determine connection mode
+            if self._d402_account and D402_AVAILABLE:
+                logger.info("💳 TraiaMCPAdapter: d402 payment protocol enabled")
+                logger.info(f"   Client account: {self._d402_account.address}")
+                logger.info(f"   Max value: {self._d402_max_value}")
+                logger.info("   Applying httpx monkey-patch for d402 hooks...")
+                # Apply d402 payment hooks patch
+                self._apply_d402_patch()
+                logger.info("   ✅ Monkey-patch applied - ALL new httpx clients will have d402 hooks")
+            elif self._auth_headers:
                 logger.info(f"TraiaMCPAdapter: Headers configured: {list(self._auth_headers.keys())}")
-                # Apply monkey patch before parent initialization
+                # Apply header injection patch
                 self._apply_httpx_patch()
             
-            # Pass clean params to parent
+            # Pass clean params to parent - this will create the MCP client
+            logger.info("   Creating parent MCPServerAdapter...")
             super().__init__(server_params_copy, *tool_names)
+            logger.info("   ✅ Parent MCPServerAdapter created")
         else:
-            # Non-dict params, no headers possible
+            # Non-dict params, no headers or payment possible
             self._auth_headers = {}
+            self._d402_account = None
+            self._d402_max_value = None
             super().__init__(server_params, *tool_names)
     
     def _apply_httpx_patch(self):
@@ -107,7 +140,56 @@ class TraiaMCPAdapter(MCPServerAdapter):
         # Apply the patch
         httpx.AsyncClient.__init__ = patched_init
         self._original_httpx_init = original_init
-        logger.debug("Applied httpx.AsyncClient monkey patch")
+        logger.debug("Applied httpx.AsyncClient monkey patch for headers")
+    
+    def _apply_d402_patch(self):
+        """Monkey-patch httpx.AsyncClient to add d402 payment hooks."""
+        if not D402_AVAILABLE:
+            raise ImportError("d402 payment hooks not available. Ensure traia_iatp.d402 is installed.")
+        
+        original_init = httpx.AsyncClient.__init__
+        d402_account = self._d402_account  # Keep variable name for compatibility
+        d402_max_value = self._d402_max_value
+        
+        @wraps(original_init)
+        def patched_init(client_self, *args, **kwargs):
+            # Log client creation
+            import traceback
+            creation_stack = ''.join(traceback.format_stack()[-4:-1])
+            logger.info(f"🔨 httpx.AsyncClient.__init__ called")
+            logger.debug(f"   Creation stack: {creation_stack}")
+            
+            # Call original init first
+            original_init(client_self, *args, **kwargs)
+            
+            # Add d402 payment hooks to the client
+            if d402_account:
+                try:
+                    hooks = d402_payment_hooks(d402_account, max_value=d402_max_value)
+                    # Merge with existing hooks if any
+                    if hasattr(client_self, 'event_hooks') and client_self.event_hooks:
+                        # Merge hooks
+                        for event_type, handlers in hooks.items():
+                            if event_type in client_self.event_hooks:
+                                client_self.event_hooks[event_type].extend(handlers)
+                            else:
+                                client_self.event_hooks[event_type] = handlers
+                    else:
+                        client_self.event_hooks = hooks
+                    
+                        logger.info(f"✅ Applied d402 payment hooks to httpx.AsyncClient")
+                        logger.info(f"   Account: {d402_account.address[:10]}...")
+                        logger.info(f"   Hooks: {list(hooks.keys())}")
+                except Exception as e:
+                    logger.error(f"❌ Failed to apply d402 hooks: {e}")
+                    import traceback
+                    logger.error(traceback.format_exc())
+                    raise
+        
+        # Apply the patch
+        httpx.AsyncClient.__init__ = patched_init
+        self._original_httpx_init = original_init
+        logger.debug("Applied httpx.AsyncClient monkey patch for d402 payments")
     
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit context manager and restore original httpx if patched."""
@@ -121,7 +203,9 @@ class TraiaMCPAdapter(MCPServerAdapter):
     
     def __enter__(self) -> List[BaseTool]:
         """Enter context manager."""
-        if self._auth_headers:
+        if self._d402_account:
+            logger.debug(f"TraiaMCPAdapter: Using d402 payment connection")
+        elif self._auth_headers:
             logger.debug(f"TraiaMCPAdapter: Using authenticated connection")
         
         return super().__enter__()
@@ -131,29 +215,45 @@ def create_mcp_adapter(
     url: str,
     headers: Optional[Dict[str, str]] = None,
     transport: str = "streamable-http",
-    tool_names: Optional[List[str]] = None
+    tool_names: Optional[List[str]] = None,
+    x402_account: Optional[Any] = None,
+    x402_max_value: Optional[int] = None
 ) -> TraiaMCPAdapter:
     """
-    Create a Traia MCP adapter with optional headers.
+    Create a Traia MCP adapter with optional headers and/or d402 payment support.
     
     Args:
         url: MCP server URL
         headers: Optional dictionary of headers to include in requests
         transport: Transport type (default: streamable-http)
         tool_names: Optional list of tool names to filter
+        x402_account: Optional eth_account.Account for d402 payment protocol
+        x402_max_value: Optional maximum payment value in base units.
+                       Typically, each MCP server uses one primary token, so this limit
+                       applies to all endpoints using that token. Set based on the token's
+                       base units (e.g., USDC with 6 decimals: $1.00 = 1_000_000).
     
     Returns:
-        TraiaMCPAdapter configured with or without headers
+        TraiaMCPAdapter configured with headers and/or d402 payments
     
     Example:
         ```python
-        # With headers
+        # With headers (authenticated)
         adapter = create_mcp_adapter(
             url="https://news-mcp.example.com/mcp",
             headers={"Authorization": "Bearer YOUR_API_KEY"}
         )
         
-        # Without headers
+        # With d402 payments
+        from eth_account import Account
+        account = Account.from_key("0x...")
+        adapter = create_mcp_adapter(
+            url="https://news-mcp.example.com/mcp",
+            x402_account=account,
+            x402_max_value=1_000_000  # $1.00 in USDC (6 decimals)
+        )
+        
+        # Without headers or payment (standard)
         adapter = create_mcp_adapter(
             url="https://news-mcp.example.com/mcp"
         )
@@ -170,6 +270,11 @@ def create_mcp_adapter(
     if headers:
         server_params["headers"] = headers
     
+    if x402_account:
+        server_params["d402_account"] = x402_account  # Use d402_ internally
+        if x402_max_value is not None:
+            server_params["d402_max_value"] = x402_max_value  # Use d402_ internally
+    
     if tool_names:
         return TraiaMCPAdapter(server_params, *tool_names)
     else:
@@ -219,6 +324,62 @@ def create_mcp_adapter_with_auth(
     return create_mcp_adapter(url, headers, transport, tool_names)
 
 
+def create_mcp_adapter_with_x402(
+    url: str,
+    account: Any,
+    max_value: Optional[int] = None,
+    transport: str = "streamable-http",
+    tool_names: Optional[List[str]] = None
+) -> TraiaMCPAdapter:
+    """
+    Create a Traia MCP adapter with d402 payment protocol support.
+    
+    Note: Function name uses "x402" for backward compatibility, but it uses
+    the d402 implementation from traia_iatp.d402.
+    
+    Args:
+        url: MCP server URL
+        account: eth_account.Account instance for signing payments
+        max_value: Optional maximum allowed payment amount in base units.
+                   Typically, each MCP server uses one primary token, so this limit
+                   applies to all endpoints using that token. Set based on the token's
+                   base units (e.g., USDC with 6 decimals: $1.00 = 1_000_000).
+        transport: Transport type (default: streamable-http)
+        tool_names: Optional list of tool names to filter
+    
+    Returns:
+        TraiaMCPAdapter configured with d402 payment hooks
+    
+    Example:
+        ```python
+        from eth_account import Account
+        
+        account = Account.from_key("0x...")
+        adapter = create_mcp_adapter_with_x402(
+            url="https://news-mcp.example.com/mcp",
+            account=account,
+            max_value=1_000_000  # $1.00 in USDC (6 decimals) - adjust for your server's token
+        )
+        
+        with adapter as tools:
+            # Tools will automatically handle 402 Payment Required responses
+            agent = Agent(tools=tools)
+        ```
+    """
+    if not D402_AVAILABLE:
+        raise ImportError(
+            "d402 payment hooks not available. Ensure traia_iatp.d402 is installed."
+        )
+    
+    return create_mcp_adapter(
+        url=url,
+        transport=transport,
+        tool_names=tool_names,
+        x402_account=account,  # Function param uses x402 for backward compat
+        x402_max_value=max_value  # Will be converted to d402 internally
+    )
+
+
 # Backwards compatibility aliases
 HeaderMCPAdapter = TraiaMCPAdapter
 create_mcp_adapter_with_headers = create_mcp_adapter
@@ -236,7 +397,7 @@ all HTTP requests when using streamable-http transport.
 Basic Usage
 -----------
 ```python
-from iatp.mcp import TraiaMCPAdapter
+from traia_iatp.mcp import TraiaMCPAdapter
 
 # Standard connection (no authentication)
 server_params = {
@@ -264,7 +425,7 @@ with TraiaMCPAdapter(server_params) as tools:
 Using Helper Functions
 ----------------------
 ```python
-from iatp.mcp import create_mcp_adapter_with_auth
+from traia_iatp.mcp import create_mcp_adapter_with_auth
 
 # Create authenticated adapter
 adapter = create_mcp_adapter_with_auth(
@@ -321,8 +482,8 @@ class AuthMiddleware(Middleware):
             token = auth[7:].strip() if auth.lower().startswith("bearer ") else None
             
             if not token:
-                # Check X-API-KEY header as alternative
-                token = request.headers.get("X-API-KEY", "")
+                # Check x-api-key header as alternative (case-insensitive)
+                token = request.headers.get("x-api-key", "")
             
             if token:
                 # Store the API key in the context state
@@ -376,6 +537,7 @@ __all__ = [
     'HeaderMCPAdapter',  # Alias for backward compatibility
     'create_mcp_adapter',
     'create_mcp_adapter_with_auth',
+    'create_mcp_adapter_with_x402',
     'create_mcp_adapter_with_headers',  # Alias
     'USAGE_GUIDE'
 ] 

traia_iatp/registry/__init__.py

@@ -1,16 +1,37 @@
-"""IATP registry module for managing utility agents and MCP servers."""
+"""IATP registry module for managing utility agents and MCP servers with lazy loading."""
 
-from .mongodb_registry import UtilityAgentRegistry, MCPServerRegistry
-from .iatp_search_api import (
-    find_utility_agent, 
-    list_utility_agents, 
-    search_utility_agents,
-    find_mcp_server,
-    list_mcp_servers,
-    search_mcp_servers,
-    get_mcp_server
-)
-from .embeddings import get_embedding_service
+from typing import TYPE_CHECKING
+
+# Type hints for IDEs and type checkers (not loaded at runtime)
+if TYPE_CHECKING:
+    from .mongodb_registry import UtilityAgentRegistry, MCPServerRegistry
+    from .iatp_search_api import (
+        find_utility_agent, 
+        list_utility_agents, 
+        search_utility_agents,
+        find_mcp_server,
+        list_mcp_servers,
+        search_mcp_servers,
+        get_mcp_server
+    )
+    from .embeddings import get_embedding_service
+
+# Lazy imports to avoid loading heavy dependencies (OpenAI, Cohere, etc.) unless needed
+_LAZY_IMPORTS = {
+    # Registry classes (write operations)
+    "UtilityAgentRegistry": ".mongodb_registry",
+    "MCPServerRegistry": ".mongodb_registry",
+    # Search API functions (lightweight, no heavy deps)
+    "find_utility_agent": ".iatp_search_api",
+    "list_utility_agents": ".iatp_search_api",
+    "search_utility_agents": ".iatp_search_api",
+    "find_mcp_server": ".iatp_search_api",
+    "list_mcp_servers": ".iatp_search_api",
+    "search_mcp_servers": ".iatp_search_api",
+    "get_mcp_server": ".iatp_search_api",
+    # Embedding service (HEAVY - OpenAI/Cohere, only load if vector search is used)
+    "get_embedding_service": ".embeddings",
+}
 
 __all__ = [
     "UtilityAgentRegistry",
@@ -24,3 +45,17 @@ __all__ = [
     "get_mcp_server",
     "get_embedding_service",
 ]
+
+
+def __getattr__(name: str):
+    """Lazy import mechanism to load modules only when accessed."""
+    if name in _LAZY_IMPORTS:
+        from importlib import import_module
+        module_path = _LAZY_IMPORTS[name]
+        module = import_module(module_path, package=__package__)
+        attr = getattr(module, name)
+        # Cache the imported attribute
+        globals()[name] = attr
+        return attr
+    
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

traia_iatp/registry/atlas_search_indexes.json

@@ -6,43 +6,70 @@
       "mappings": {
         "dynamic": false,
         "fields": {
-          "name": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
-          "description": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
-          "tags": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
           "capabilities": {
-            "type": "string",
-            "analyzer": "lucene.standard"
+            "analyzer": "lucene.standard",
+            "type": "string"
           },
-          "search_text": {
-            "type": "string",
-            "analyzer": "lucene.standard"
+          "description": {
+            "analyzer": "lucene.standard",
+            "type": "string"
           },
+          "name": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
+            }
+          ],
+          "search_text": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
+            }
+          ],
           "skills": {
-            "type": "document",
             "fields": {
-              "name": {
-                "type": "string",
-                "analyzer": "lucene.standard"
-              },
               "description": {
-                "type": "string",
-                "analyzer": "lucene.standard"
+                "analyzer": "lucene.standard",
+                "type": "string"
+              },
+              "name": {
+                "analyzer": "lucene.standard",
+                "type": "string"
               },
               "tags": {
-                "type": "string",
-                "analyzer": "lucene.standard"
+                "analyzer": "lucene.standard",
+                "type": "string"
               }
+            },
+            "type": "document"
+          },
+          "tags": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
             }
-          }
+          ]
         }
       }
     }
@@ -54,43 +81,70 @@
       "mappings": {
         "dynamic": false,
         "fields": {
-          "name": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
-          "description": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
-          "tags": {
-            "type": "string",
-            "analyzer": "lucene.standard"
-          },
           "capabilities": {
-            "type": "string",
-            "analyzer": "lucene.standard"
+            "analyzer": "lucene.standard",
+            "type": "string"
           },
-          "search_text": {
-            "type": "string",
-            "analyzer": "lucene.standard"
+          "description": {
+            "analyzer": "lucene.standard",
+            "type": "string"
           },
+          "name": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
+            }
+          ],
+          "search_text": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
+            }
+          ],
           "skills": {
-            "type": "document",
             "fields": {
-              "name": {
-                "type": "string",
-                "analyzer": "lucene.standard"
-              },
               "description": {
-                "type": "string",
-                "analyzer": "lucene.standard"
+                "analyzer": "lucene.standard",
+                "type": "string"
+              },
+              "name": {
+                "analyzer": "lucene.standard",
+                "type": "string"
               },
               "tags": {
-                "type": "string",
-                "analyzer": "lucene.standard"
+                "analyzer": "lucene.standard",
+                "type": "string"
               }
+            },
+            "type": "document"
+          },
+          "tags": [
+            {
+              "analyzer": "lucene.standard",
+              "type": "string"
+            },
+            {
+              "foldDiacritics": true,
+              "maxGrams": 15,
+              "minGrams": 2,
+              "tokenization": "edgeGram",
+              "type": "autocomplete"
             }
-          }
+          ]
         }
       }
     }