traia-iatp 0.1.29__py3-none-any.whl

traia_iatp/mcp/d402_mcp_tool_adapter.py

@@ -0,0 +1,361 @@
+#!/usr/bin/env python
+"""
+D402 MCP Tool Adapter
+
+A simple adapter for using d402-enabled MCP servers with CrewAI.
+This adapter avoids the complexity of persistent SSE connections and background tasks.
+
+Instead, it:
+1. Lists available tools from the MCP server
+2. Creates CrewAI BaseTool wrappers for each MCP tool
+3. Each tool uses httpx with d402 payment hooks for requests
+4. No persistent connections - simple request/response pattern
+
+This is more reliable than MCPServerAdapter for d402 payment scenarios.
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional, Type
+from pydantic import BaseModel, Field
+from crewai.tools import BaseTool
+import httpx
+
+from traia_iatp.d402.clients.httpx import d402_payment_hooks
+
+logger = logging.getLogger(__name__)
+
+
+class D402MCPTool(BaseTool):
+    """
+    CrewAI tool wrapper for a single MCP tool with d402 payment support.
+    
+    Each instance represents one MCP tool and handles d402 payments automatically.
+    """
+    
+    name: str = "mcp_tool"
+    description: str = "MCP tool with d402 payment"
+    mcp_server_url: str = ""
+    mcp_tool_name: str = ""
+    mcp_session_id: str = ""
+    d402_operator_account: Optional[Any] = None  # Operator account (EOA) for signing
+    d402_wallet_address: Optional[str] = None  # IATPWallet contract address
+    input_schema: Dict[str, Any] = Field(default_factory=dict)
+    
+    class Config:
+        arbitrary_types_allowed = True
+    
+    def __init__(self, **data):
+        """Initialize with custom data."""
+        super().__init__(**data)
+    
+    def _run(self, **kwargs) -> Dict[str, Any]:
+        """Synchronous wrapper that runs async _arun."""
+        import asyncio
+        try:
+            # Run the async method in a new event loop
+            return asyncio.run(self._arun(**kwargs))
+        except RuntimeError:
+            # If we're already in an event loop, use run_until_complete
+            loop = asyncio.get_event_loop()
+            return loop.run_until_complete(self._arun(**kwargs))
+    
+    async def _arun(self, **kwargs) -> Dict[str, Any]:
+        """
+        Execute the MCP tool with d402 payment support.
+        
+        This method:
+        1. Creates httpx client with d402 hooks
+        2. Makes MCP tools/call request
+        3. d402 hooks automatically handle any 402 responses
+        4. Returns the result
+        """
+        
+        # Create httpx client with d402 payment hooks
+        async with httpx.AsyncClient(base_url=self.mcp_server_url, timeout=60.0) as client:
+            # Attach d402 payment hooks
+            client.event_hooks = d402_payment_hooks(
+                self.d402_operator_account,
+                wallet_address=self.d402_wallet_address
+            )
+            
+            # Create MCP tools/call request
+            mcp_request = {
+                "jsonrpc": "2.0",
+                "id": 1,
+                "method": "tools/call",
+                "params": {
+                    "name": self.mcp_tool_name,
+                    "arguments": kwargs
+                }
+            }
+            
+            try:
+                # Make request - d402 hooks will handle any 402 responses automatically
+                response = await client.post(
+                    "/",
+                    json=mcp_request,
+                    headers={
+                        "Content-Type": "application/json",
+                        "Accept": "application/json, text/event-stream",
+                        "mcp-session-id": self.mcp_session_id
+                    },
+                    follow_redirects=True
+                )
+                
+                # Read response content
+                content = await response.aread()
+                content_str = content.decode() if content else ""
+                
+                # Parse SSE format if present
+                if "data:" in content_str:
+                    sse_lines = content_str.strip().split('\n')
+                    for line in sse_lines:
+                        if line.startswith("data:"):
+                            content_str = line[5:].strip()
+                            break
+                
+                # Parse JSON response
+                response_data = json.loads(content_str)
+                
+                # Extract result
+                if "result" in response_data:
+                    result = response_data["result"]
+                    # Return structured content if available, otherwise full result
+                    if isinstance(result, dict) and "structuredContent" in result:
+                        return result["structuredContent"].get("result", result)
+                    return result
+                elif "error" in response_data:
+                    return {"error": response_data["error"]}
+                else:
+                    return response_data
+                    
+            except Exception as e:
+                logger.error(f"Error calling MCP tool {self.mcp_tool_name}: {e}")
+                return {"error": str(e)}
+
+
+class D402MCPToolAdapter:
+    """
+    Adapter for using d402-enabled MCP servers with CrewAI.
+    
+    This adapter is simpler and more reliable than MCPServerAdapter for d402:
+    - No persistent SSE connections
+    - No background tasks
+    - Direct request/response with d402 hooks
+    
+    Usage:
+        ```python
+        from eth_account import Account
+        from traia_iatp.mcp import D402MCPToolAdapter
+        
+        account = Account.from_key("0x...")
+        adapter = D402MCPToolAdapter(
+            url="http://localhost:8000/mcp",
+            account=account
+        )
+        
+        with adapter as tools:
+            agent = Agent(
+                role="Analyst",
+                goal="Analyze data",
+                tools=tools
+            )
+        ```
+    """
+    
+    def __init__(
+        self,
+        url: str,
+        account: Any,  # Operator account (EOA) for signing
+        wallet_address: str = None,  # IATPWallet contract address
+        max_value: Optional[int] = None
+    ):
+        """
+        Initialize the d402 MCP adapter.
+        
+        Args:
+            url: MCP server URL (e.g., "http://localhost:8000/mcp")
+            account: Operator account (EOA) with private key for signing payments
+            wallet_address: IATPWallet contract address (if None, uses account.address for testing)
+            max_value: Optional maximum payment value in base units
+        """
+        self.url = url
+        self.account = account  # Operator EOA
+        self.wallet_address = wallet_address or account.address  # IATPWallet or EOA for testing
+        self.max_value = max_value
+        self.tools: List[BaseTool] = []
+        self.session_id: Optional[str] = None
+    
+    async def _initialize_session(self) -> str:
+        """Initialize MCP session and return session ID."""
+        async with httpx.AsyncClient(base_url=self.url, timeout=30.0) as client:
+            init_request = {
+                "jsonrpc": "2.0",
+                "id": "init",
+                "method": "initialize",
+                "params": {
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {},
+                    "clientInfo": {"name": "d402-crewai-client", "version": "1.0"}
+                }
+            }
+            
+            response = await client.post(
+                "/",
+                json=init_request,
+                headers={
+                    "Content-Type": "application/json",
+                    "Accept": "application/json, text/event-stream"
+                },
+                follow_redirects=True
+            )
+            
+            session_id = response.headers.get("mcp-session-id")
+            if not session_id:
+                raise RuntimeError("Failed to establish MCP session")
+            
+            return session_id
+    
+    async def _list_tools(self) -> List[Dict[str, Any]]:
+        """List available tools from the MCP server."""
+        async with httpx.AsyncClient(base_url=self.url, timeout=30.0) as client:
+            mcp_request = {
+                "jsonrpc": "2.0",
+                "id": 1,
+                "method": "tools/list",
+                "params": {}
+            }
+            
+            response = await client.post(
+                "/",
+                json=mcp_request,
+                headers={
+                    "Content-Type": "application/json",
+                    "Accept": "application/json, text/event-stream",
+                    "mcp-session-id": self.session_id
+                },
+                follow_redirects=True
+            )
+            
+            content = await response.aread()
+            content_str = content.decode() if content else ""
+            
+            # Parse SSE format if present
+            if "data:" in content_str:
+                sse_lines = content_str.strip().split('\n')
+                for line in sse_lines:
+                    if line.startswith("data:"):
+                        content_str = line[5:].strip()
+                        break
+            
+            response_data = json.loads(content_str)
+            
+            if "result" in response_data and "tools" in response_data["result"]:
+                return response_data["result"]["tools"]
+            return []
+    
+    def __enter__(self) -> List[BaseTool]:
+        """
+        Enter context manager - set up session and create tool wrappers.
+        
+        Returns:
+            List of CrewAI BaseTool objects for each MCP tool
+        """
+        import asyncio
+        
+        # Initialize session
+        logger.info(f"🔌 Initializing MCP session with {self.url}")
+        self.session_id = asyncio.run(self._initialize_session())
+        logger.info(f"✅ Session established: {self.session_id}")
+        
+        # List available tools
+        logger.info("📋 Listing available tools...")
+        mcp_tools = asyncio.run(self._list_tools())
+        logger.info(f"✅ Found {len(mcp_tools)} tools")
+        
+        # Create CrewAI tool wrappers for each MCP tool
+        self.tools = []
+        for mcp_tool in mcp_tools:
+            tool_name = mcp_tool.get("name", "unknown")
+            tool_description = mcp_tool.get("description", f"MCP tool: {tool_name}")
+            
+            # Create tool instance with all necessary data
+            tool_instance = D402MCPTool(
+                name=tool_name,
+                description=tool_description,
+                mcp_server_url=self.url,
+                mcp_tool_name=tool_name,
+                mcp_session_id=self.session_id,
+                d402_operator_account=self.account,  # Operator EOA for signing
+                d402_wallet_address=self.wallet_address,  # IATPWallet contract (or EOA for testing)
+                input_schema=mcp_tool.get("inputSchema", {})
+            )
+            
+            self.tools.append(tool_instance)
+        
+        logger.info(f"✅ Created {len(self.tools)} CrewAI tool wrappers")
+        return self.tools
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager - cleanup."""
+        logger.info("🔌 Closing D402MCPToolAdapter")
+        return False
+
+
+def create_d402_mcp_adapter(
+    url: str,
+    account: Any,  # Operator account (EOA) for signing
+    wallet_address: str = None,  # IATPWallet contract address
+    max_value: Optional[int] = None
+) -> D402MCPToolAdapter:
+    """
+    Create a d402 MCP adapter for CrewAI.
+    
+    Args:
+        url: MCP server URL
+        account: Operator account (EOA) for signing payments
+        wallet_address: IATPWallet contract address (if None, uses account.address for testing)
+        max_value: Optional maximum payment value in base units
+        
+    Returns:
+        D402MCPToolAdapter instance
+        
+    Example:
+        ```python
+        from eth_account import Account
+        from traia_iatp.mcp import create_d402_mcp_adapter
+        
+        # For testing (uses mock wallet address)
+        operator_account = Account.from_key("0x...")
+        mock_wallet = "0x70997970C51812dc3A010C7d01b50e0d17dc79C8"  # Different from operator
+        adapter = create_d402_mcp_adapter(
+            url="http://localhost:8000/mcp",
+            account=operator_account,
+            wallet_address=mock_wallet,
+            max_value=1_000_000  # $1.00 in USDC
+        )
+        
+        # For production (with deployed IATPWallet contract)
+        operator_account = Account.from_key("0x...")  # Operator EOA
+        wallet_address = "0x..."  # Deployed IATPWallet contract
+        adapter = create_d402_mcp_adapter(
+            url="http://localhost:8000/mcp",
+            account=operator_account,
+            wallet_address=wallet_address,
+            max_value=1_000_000
+        )
+        
+        with adapter as tools:
+            agent = Agent(
+                role="Analyst",
+                goal="Analyze data",
+                tools=tools
+            )
+        ```
+    """
+    return D402MCPToolAdapter(url, account, wallet_address, max_value)
+
+
+__all__ = ["D402MCPTool", "D402MCPToolAdapter", "create_d402_mcp_adapter"]
+