atp-protocol 1.0.0__py3-none-any.whl

atp/__init__.py

@@ -0,0 +1,11 @@
+"""ATP Protocol package."""
+
+from atp.middleware import (
+    ATPSettlementMiddleware,
+    create_settlement_middleware,
+)
+
+__all__ = [
+    "ATPSettlementMiddleware",
+    "create_settlement_middleware",
+]

atp/config.py

@@ -0,0 +1,77 @@
+"""
+Central configuration for the ATP Gateway.
+
+Keep all env parsing + constants here so the rest of the code can be imported cleanly.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def _bool_env(name: str, default: bool = False) -> bool:
+    v = os.getenv(name)
+    if v is None:
+        return default
+    return v.strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
+def _float_env(name: str) -> Optional[float]:
+    v = os.getenv(name)
+    if v is None:
+        return None
+    v = v.strip()
+    if not v:
+        return None
+    return float(v)
+
+
+SWARMS_API_KEY = os.getenv("SWARMS_API_KEY")
+SWARMS_API_URL = os.getenv(
+    "SWARMS_API_URL", "https://api.swarms.world/v1/agent/completions"
+)
+
+SOLANA_RPC_URL = os.getenv(
+    "SOLANA_RPC_URL", "https://api.mainnet-beta.solana.com"
+)
+AGENT_TREASURY_PUBKEY = os.getenv("AGENT_TREASURY_PUBKEY")
+
+JOB_TTL_SECONDS = int(os.getenv("JOB_TTL_SECONDS", "600"))
+
+# Swarms Treasury for settlement fees
+SWARMS_TREASURY_PUBKEY = os.getenv(
+    "SWARMS_TREASURY_PUBKEY",
+    "7MaX4muAn8ZQREJxnupm8sgokwFHujgrGfH9Qn81BuEV",
+)
+SETTLEMENT_FEE_PERCENT = float(
+    os.getenv("SETTLEMENT_FEE_PERCENT", "0.05")
+)
+
+# USDC Token Configuration (Solana Mainnet)
+USDC_MINT_ADDRESS = os.getenv(
+    "USDC_MINT_ADDRESS",
+    "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
+)
+USDC_DECIMALS = int(os.getenv("USDC_DECIMALS", "6"))
+
+SUPPORTED_PAYMENT_TOKENS = ["SOL", "USDC"]
+
+# Agent pricing configuration (host-provided)
+INPUT_COST_PER_MILLION_USD = _float_env("INPUT_COST_PER_MILLION_USD")
+OUTPUT_COST_PER_MILLION_USD = _float_env(
+    "OUTPUT_COST_PER_MILLION_USD"
+)
+
+# Verbose Solana debug logging (server-side). Enable temporarily for diagnosing RPC type mismatches.
+# WARNING: do not enable permanently in high-traffic production environments.
+ATP_SOLANA_DEBUG = _bool_env("ATP_SOLANA_DEBUG", default=False)
+
+# Settlement Service URL
+ATP_SETTLEMENT_URL = os.getenv(
+    "ATP_SETTLEMENT_URL", "http://localhost:8001"
+)

atp/middleware.py

@@ -0,0 +1,343 @@
+"""
+FastAPI middleware for ATP settlement on any endpoint.
+
+This middleware enables automatic payment deduction from Solana wallets
+based on token usage (input/output tokens) for any configured endpoint.
+
+The middleware delegates all settlement logic to the ATP Settlement Service,
+ensuring immutable and centralized settlement operations.
+
+The middleware accepts wallet private keys directly via headers, making it
+simple to use without requiring API key management. Users can add their
+own API key handling layer if needed.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Callable, Dict, List, Optional, Set
+
+from fastapi import HTTPException, Request, Response
+from loguru import logger
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+
+from atp import config
+from atp.schemas import PaymentToken
+from atp.settlement_client import SettlementServiceClient
+
+
+class ATPSettlementMiddleware(BaseHTTPMiddleware):
+    """
+    FastAPI middleware that automatically deducts payment from Solana wallets
+    based on token usage for configured endpoints.
+
+    The middleware delegates all settlement logic to the ATP Settlement Service,
+    ensuring immutable and centralized settlement operations.
+
+    The middleware accepts wallet private keys directly via headers, making it
+    simple to use. Users can add their own API key handling layer if needed.
+
+    Payments are split automatically:
+    - Treasury (SWARMS_TREASURY_PUBKEY) receives the processing fee
+    - Recipient (endpoint host) receives the remainder
+
+    Usage:
+        app.add_middleware(
+            ATPSettlementMiddleware,
+            allowed_endpoints=["/v1/chat", "/v1/completions"],
+            input_cost_per_million_usd=10.0,
+            output_cost_per_million_usd=30.0,
+            wallet_private_key_header="x-wallet-private-key",
+            payment_token=PaymentToken.SOL,
+            recipient_pubkey="YourPublicKeyHere",  # Required: endpoint host receives main payment
+            # settlement_service_url is optional - uses ATP_SETTLEMENT_URL env var by default
+        )
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        *,
+        allowed_endpoints: List[str],
+        input_cost_per_million_usd: float,
+        output_cost_per_million_usd: float,
+        wallet_private_key_header: str = "x-wallet-private-key",
+        payment_token: PaymentToken = PaymentToken.SOL,
+        recipient_pubkey: Optional[str] = None,
+        skip_preflight: bool = False,
+        commitment: str = "confirmed",
+        usage_response_key: str = "usage",
+        include_usage_in_response: bool = True,
+        require_wallet: bool = True,
+        settlement_service_url: Optional[str] = None,
+    ):
+        """
+        Initialize the ATP settlement middleware.
+
+        The middleware delegates all settlement logic to the ATP Settlement Service.
+        All settlement operations are handled by the immutable settlement service.
+
+        Args:
+            app: The ASGI application.
+            allowed_endpoints: List of endpoint paths to apply settlement to (e.g., ["/v1/chat"]).
+                Supports path patterns - exact matches only.
+            input_cost_per_million_usd: Cost per million input tokens in USD.
+            output_cost_per_million_usd: Cost per million output tokens in USD.
+            wallet_private_key_header: HTTP header name containing the wallet private key
+                (default: "x-wallet-private-key"). The private key should be in JSON array
+                format (e.g., "[1,2,3,...]") or base58 string format.
+            payment_token: Token to use for payment (SOL or USDC).
+            recipient_pubkey: Solana public key of the recipient wallet (the endpoint host).
+                This wallet receives the main payment (after processing fee). Required.
+            skip_preflight: Whether to skip preflight simulation for Solana transactions.
+            commitment: Solana commitment level (processed|confirmed|finalized).
+            usage_response_key: Key in response JSON where usage data is located (default: "usage").
+            include_usage_in_response: Whether to add usage/cost info to the response.
+            require_wallet: Whether to require wallet private key (if False, skips settlement when missing).
+            settlement_service_url: Base URL of the settlement service. If not provided, uses
+                ATP_SETTLEMENT_URL environment variable (default: http://localhost:8001).
+                The middleware always uses the settlement service for all settlement operations.
+        """
+        super().__init__(app)
+        self.allowed_endpoints: Set[str] = set(allowed_endpoints)
+        self.input_cost_per_million_usd = input_cost_per_million_usd
+        self.output_cost_per_million_usd = output_cost_per_million_usd
+        self.wallet_private_key_header = (
+            wallet_private_key_header.lower()
+        )
+        self.payment_token = payment_token
+        # Recipient pubkey - configurable, the endpoint host receives the main payment
+        self._recipient_pubkey = recipient_pubkey
+        if not self._recipient_pubkey:
+            raise ValueError("recipient_pubkey must be provided")
+        # Treasury pubkey - always uses SWARMS_TREASURY_PUBKEY for processing fees
+        self._treasury_pubkey = config.SWARMS_TREASURY_PUBKEY
+        if not self._treasury_pubkey:
+            raise ValueError(
+                "SWARMS_TREASURY_PUBKEY must be set in configuration"
+            )
+        self.skip_preflight = skip_preflight
+        self.commitment = commitment
+        self.usage_response_key = usage_response_key
+        self.include_usage_in_response = include_usage_in_response
+        self.require_wallet = require_wallet
+        # Always use settlement service - initialize client with config value or provided URL
+        service_url = (
+            settlement_service_url or config.ATP_SETTLEMENT_URL
+        )
+        self.settlement_service_client = SettlementServiceClient(
+            base_url=service_url
+        )
+
+    def _should_process(self, path: str) -> bool:
+        """Check if the request path should be processed by this middleware."""
+        return path in self.allowed_endpoints
+
+    def _extract_wallet_private_key(
+        self, request: Request
+    ) -> Optional[str]:
+        """Extract wallet private key from request headers."""
+        return request.headers.get(self.wallet_private_key_header)
+
+    async def _extract_usage_from_response(
+        self, response_body: bytes
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Extract usage information from response body.
+
+        Tries multiple strategies:
+        1. Look for usage data at the configured usage_response_key
+        2. Check if the entire response contains usage-like keys
+        3. Try nested structures (usage.usage, meta.usage, etc.)
+
+        The usage data is then sent to the settlement service for parsing,
+        so we just need to extract the raw usage object.
+        """
+        try:
+            body_str = response_body.decode("utf-8")
+            if not body_str.strip():
+                return None
+            data = json.loads(body_str)
+
+            # Strategy 1: Try the configured usage key first
+            usage = data.get(self.usage_response_key)
+            if usage and isinstance(usage, dict):
+                return usage
+
+            # Strategy 2: Check if the entire response is usage-like
+            if isinstance(data, dict):
+                # Check for common usage keys at top level
+                usage_keys = [
+                    "input_tokens",
+                    "output_tokens",
+                    "prompt_tokens",
+                    "completion_tokens",
+                    "total_tokens",
+                    "tokens",
+                    "promptTokenCount",
+                    "candidatesTokenCount",
+                    "totalTokenCount",
+                ]
+                if any(key in data for key in usage_keys):
+                    return data
+
+            # Strategy 3: Try nested structures
+            # Check for usage nested in common locations
+            for nested_key in [
+                "usage",
+                "token_usage",
+                "tokens",
+                "statistics",
+                "meta",
+            ]:
+                if nested_key in data and isinstance(
+                    data[nested_key], dict
+                ):
+                    nested_usage = data[nested_key]
+                    # Check if it looks like usage data
+                    if any(
+                        key in nested_usage
+                        for key in [
+                            "input_tokens",
+                            "output_tokens",
+                            "prompt_tokens",
+                            "completion_tokens",
+                            "tokens",
+                        ]
+                    ):
+                        return nested_usage
+
+            return None
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.debug(
+                f"Failed to parse response body for usage: {e}"
+            )
+            return None
+
+    async def dispatch(
+        self, request: Request, call_next: Callable
+    ) -> Response:
+        """Process the request and apply settlement if applicable."""
+        path = request.url.path
+
+        # Skip if not in allowed endpoints
+        if not self._should_process(path):
+            return await call_next(request)
+
+        # Extract wallet private key
+        private_key = self._extract_wallet_private_key(request)
+        if not private_key:
+            if self.require_wallet:
+                raise HTTPException(
+                    status_code=401,
+                    detail=f"Missing wallet private key in header: {self.wallet_private_key_header}",
+                )
+            # If wallet not required, skip settlement
+            return await call_next(request)
+
+        # Execute the endpoint
+        response = await call_next(request)
+
+        # Only process successful responses
+        if response.status_code >= 400:
+            return response
+
+        # Extract usage from response
+        response_body = b""
+        async for chunk in response.body_iterator:
+            response_body += chunk
+
+        usage = await self._extract_usage_from_response(response_body)
+
+        if not usage:
+            logger.warning(
+                f"No usage data found in response for {path}. Response keys: {list(json.loads(response_body.decode('utf-8')).keys()) if response_body else 'empty'}"
+            )
+            # Return original response if no usage found
+            return Response(
+                content=response_body,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                media_type=response.media_type,
+            )
+
+        # Calculate and deduct payment via settlement service
+        try:
+            payment_result = await self.settlement_service_client.settle(
+                private_key=private_key,
+                usage=usage,
+                input_cost_per_million_usd=self.input_cost_per_million_usd,
+                output_cost_per_million_usd=self.output_cost_per_million_usd,
+                recipient_pubkey=self._recipient_pubkey,
+                payment_token=self.payment_token.value,
+                treasury_pubkey=self._treasury_pubkey,
+                skip_preflight=self.skip_preflight,
+                commitment=self.commitment,
+            )
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Settlement error: {e}", exc_info=True)
+            raise HTTPException(
+                status_code=500,
+                detail=f"Settlement failed: {str(e)}",
+            )
+
+        # Modify response to include usage/payment info if requested
+        if self.include_usage_in_response:
+            try:
+                response_data = json.loads(
+                    response_body.decode("utf-8")
+                )
+                response_data["atp_settlement"] = payment_result
+                response_data["atp_usage"] = usage
+                response_body = json.dumps(response_data).encode(
+                    "utf-8"
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Failed to add settlement info to response: {e}"
+                )
+
+        return Response(
+            content=response_body,
+            status_code=response.status_code,
+            headers=dict(response.headers),
+            media_type=response.media_type,
+        )
+
+
+def create_settlement_middleware(
+    allowed_endpoints: List[str],
+    input_cost_per_million_usd: float,
+    output_cost_per_million_usd: float,
+    **kwargs: Any,
+) -> type[ATPSettlementMiddleware]:
+    """
+    Factory function to create a configured ATP settlement middleware.
+
+    Example:
+        middleware = create_settlement_middleware(
+            allowed_endpoints=["/v1/chat", "/v1/completions"],
+            input_cost_per_million_usd=10.0,
+            output_cost_per_million_usd=30.0,
+            wallet_private_key_header="x-wallet-private-key",
+            recipient_pubkey="YourPublicKeyHere",  # Optional: defaults to SWARMS_TREASURY_PUBKEY
+        )
+        app.add_middleware(middleware)
+    """
+    return type(
+        "ConfiguredATPSettlementMiddleware",
+        (ATPSettlementMiddleware,),
+        {
+            "__init__": lambda self, app: ATPSettlementMiddleware.__init__(
+                self,
+                app,
+                allowed_endpoints=allowed_endpoints,
+                input_cost_per_million_usd=input_cost_per_million_usd,
+                output_cost_per_million_usd=output_cost_per_million_usd,
+                **kwargs,
+            )
+        },
+    )

atp/schemas.py

@@ -0,0 +1,171 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+from swarms.schemas.mcp_schemas import (
+    MCPConnection,
+    MultipleMCPConnections,
+)
+
+
+class AgentSpec(BaseModel):
+    agent_name: Optional[str] = Field(
+        # default=None,
+        description="The unique name assigned to the agent, which identifies its role and functionality within the swarm.",
+    )
+    description: Optional[str] = Field(
+        default=None,
+        description="A detailed explanation of the agent's purpose, capabilities, and any specific tasks it is designed to perform.",
+    )
+    system_prompt: Optional[str] = Field(
+        default=None,
+        description="The initial instruction or context provided to the agent, guiding its behavior and responses during execution.",
+    )
+    model_name: Optional[str] = Field(
+        default="gpt-4.1",
+        description="The name of the AI model that the agent will utilize for processing tasks and generating outputs. For example: gpt-4o, gpt-4o-mini, openai/o3-mini",
+    )
+    auto_generate_prompt: Optional[bool] = Field(
+        default=False,
+        description="A flag indicating whether the agent should automatically create prompts based on the task requirements.",
+    )
+    max_tokens: Optional[int] = Field(
+        default=8192,
+        description="The maximum number of tokens that the agent is allowed to generate in its responses, limiting output length.",
+    )
+    temperature: Optional[float] = Field(
+        default=0.5,
+        description="A parameter that controls the randomness of the agent's output; lower values result in more deterministic responses.",
+    )
+    role: Optional[str] = Field(
+        default="worker",
+        description="The designated role of the agent within the swarm, which influences its behavior and interaction with other agents.",
+    )
+    max_loops: Optional[int] = Field(
+        default=1,
+        description="The maximum number of times the agent is allowed to repeat its task, enabling iterative processing if necessary.",
+    )
+    tools_list_dictionary: Optional[List[Dict[Any, Any]]] = Field(
+        default=None,
+        description="A dictionary of tools that the agent can use to complete its task.",
+    )
+    mcp_url: Optional[str] = Field(
+        default=None,
+        description="The URL of the MCP server that the agent can use to complete its task.",
+    )
+    streaming_on: Optional[bool] = Field(
+        default=False,
+        description="A flag indicating whether the agent should stream its output.",
+    )
+    llm_args: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Additional arguments to pass to the LLM such as top_p, frequency_penalty, presence_penalty, etc.",
+    )
+    dynamic_temperature_enabled: Optional[bool] = Field(
+        default=True,
+        description="A flag indicating whether the agent should dynamically adjust its temperature based on the task.",
+    )
+
+    mcp_config: Optional[MCPConnection] = Field(
+        default=None,
+        description="The MCP connection to use for the agent.",
+    )
+
+    mcp_configs: Optional[MultipleMCPConnections] = Field(
+        default=None,
+        description="The MCP connections to use for the agent. This is a list of MCP connections. Includes multiple MCP connections.",
+    )
+
+    tool_call_summary: Optional[bool] = Field(
+        default=True,
+        description="A parameter enabling an agent to summarize tool calls.",
+    )
+
+    reasoning_effort: Optional[str] = Field(
+        default=None,
+        description="The effort to put into reasoning.",
+    )
+
+    thinking_tokens: Optional[int] = Field(
+        default=None,
+        description="The number of tokens to use for thinking.",
+    )
+
+    reasoning_enabled: Optional[bool] = Field(
+        default=False,
+        description="A parameter enabling an agent to use reasoning.",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class PaymentToken(str, Enum):
+    """Supported payment tokens on Solana."""
+
+    SOL = "SOL"
+    USDC = "USDC"
+
+
+class AgentTask(BaseModel):
+    """Complete agent task request requiring full agent specification."""
+
+    agent_config: AgentSpec = Field(
+        ...,
+        description="Complete agent configuration specification matching the Swarms API AgentSpec schema",
+    )
+    task: str = Field(
+        ...,
+        description="The task or query to execute",
+        example="Analyze the latest SOL/USDC liquidity pool data and provide trading recommendations.",
+    )
+    user_wallet: str = Field(
+        ...,
+        description="The Solana public key of the sender for payment verification",
+    )
+    payment_token: PaymentToken = Field(
+        default=PaymentToken.SOL,
+        description="Payment token to use for settlement (SOL or USDC)",
+    )
+    history: Optional[Union[Dict[Any, Any], List[Dict[str, str]]]] = (
+        Field(
+            default=None,
+            description="Optional conversation history for context",
+        )
+    )
+    img: Optional[str] = Field(
+        default=None,
+        description="Optional image URL for vision tasks",
+    )
+    imgs: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of image URLs for vision tasks",
+    )
+
+
+class SettleTrade(BaseModel):
+    """Settlement request that asks the facilitator to sign+send the payment tx.
+
+    WARNING: This is custodial-like behavior. The private key is used in-memory only
+    for the duration of this request and is not persisted.
+    """
+
+    job_id: str = Field(
+        ..., description="Job ID from the trade creation response"
+    )
+    private_key: str = Field(
+        ...,
+        description=(
+            "Payer private key encoded as a string. Supported formats:\n"
+            "- Base58 keypair (common Solana secret key string)\n"
+            "- JSON array of ints (e.g. '[12,34,...]')"
+        ),
+    )
+    skip_preflight: bool = Field(
+        default=False,
+        description="Whether to skip preflight simulation",
+    )
+    commitment: str = Field(
+        default="confirmed",
+        description="Confirmation level to wait for (processed|confirmed|finalized)",
+    )