Smart Order Routing for AI Trading Agents

from dataclasses import dataclass
from typing import List, Tuple
import numpy as np

# Model for a single AMM route
@dataclass
class Route:
    route_id: str
    fee_bps: float    # total fees in basis points
    reserve_in: float  # effective input reserve depth
    reserve_out: float # effective output reserve depth
    curve_type: str   # 'cpamm', 'stableswap', 'clmm_tick'

    def output(self, amount_in: float) -> float:
        """Compute output tokens for a given input allocation."""
        net_in = amount_in * (1 - self.fee_bps / 10000)
        if self.curve_type == 'cpamm':
            return (self.reserve_out * net_in) / (self.reserve_in + net_in)
        elif self.curve_type == 'stableswap':
            # Simplified stableswap: much flatter than cpamm
            rate = self.reserve_out / self.reserve_in
            impact = 0.5 * (net_in / self.reserve_in) ** 2
            return net_in * rate * (1 - impact)
        else:
            # CLMM: treat as cpamm for simplified model
            return (self.reserve_out * net_in) / (self.reserve_in + net_in)

    def marginal_rate(self, amount_in: float, eps: float = 1.0) -> float:
        """Derivative of output w.r.t. input at current allocation."""
        return (self.output(amount_in + eps) - self.output(amount_in)) / eps


def optimal_split(routes: List[Route], total_amount: float, iterations: int = 200) -> List[float]:
    """
    Find optimal allocation across routes by equalizing marginal rates.
    Returns a list of allocations summing to total_amount.
    """
    n = len(routes)
    # Start with equal allocation
    allocs = np.full(n, total_amount / n)

    for _ in range(iterations):
        rates = np.array([r.marginal_rate(a) for r, a in zip(routes, allocs)])
        best_rate_idx = np.argmax(rates)
        worst_rate_idx = np.argmin(rates)
        if rates[best_rate_idx] - rates[worst_rate_idx] < 1e-6:
            break
        # Shift 1% of total from worst to best
        shift = total_amount * 0.01
        shift = min(shift, allocs[worst_rate_idx] * 0.5)
        allocs[best_rate_idx] += shift
        allocs[worst_rate_idx] -= shift

    # Prune allocations below 0.5% of total (dust prevention)
    min_alloc = total_amount * 0.005
    pruned = np.where(allocs < min_alloc, 0, allocs)
    # Renormalize
    if pruned.sum() > 0:
        pruned = pruned * (total_amount / pruned.sum())

    return pruned.tolist()

import numpy as np
from typing import Optional

class DynamicSlippageCalculator:
    """
    Calculates optimal slippage tolerance based on:
    - Estimated price impact of the trade
    - Current market volatility (realized vol over recent blocks)
    - Whether MEV protection is active
    - Historical revert rate feedback
    """

    BASE_SLIPPAGE_BPS = 30      # 0.30% base tolerance
    MEV_PROTECTION_DISCOUNT = 0.5 # Halve MEV buffer when protected
    REVERT_ADJUST_FACTOR = 1.2   # Expand by 20% on each revert
    MIN_SLIPPAGE_BPS = 10        # 0.10% floor
    MAX_SLIPPAGE_BPS = 200       # 2.00% ceiling

    def __init__(self):
        self._revert_history: list[bool] = []
        self._vol_window: list[float] = []  # recent block-level price changes

    def update_vol(self, block_price_change_bps: float) -> None:
        self._vol_window.append(abs(block_price_change_bps))
        if len(self._vol_window) > 50:
            self._vol_window.pop(0)

    def record_result(self, reverted: bool) -> None:
        self._revert_history.append(reverted)
        if len(self._revert_history) > 100:
            self._revert_history.pop(0)

    def _revert_rate(self) -> float:
        if not self._revert_history:
            return 0.0
        return sum(self._revert_history) / len(self._revert_history)

    def _vol_estimate_bps(self) -> float:
        if not self._vol_window:
            return 15  # 0.15% default
        return np.percentile(self._vol_window, 75)

    def calculate(
        self,
        price_impact_bps: float,
        mev_protected: bool = True,
        urgency: float = 1.0,  # 1.0 = normal, 2.0 = urgent
    ) -> float:
        """Return recommended slippage tolerance in basis points."""
        vol_bps = self._vol_estimate_bps()
        revert_rate = self._revert_rate()

        # Base: price impact + 1-block volatility buffer
        slippage = price_impact_bps + vol_bps + self.BASE_SLIPPAGE_BPS

        # MEV sandwich buffer (reduced if protected)
        mev_buffer = 30 * (1 - self.MEV_PROTECTION_DISCOUNT * mev_protected)
        slippage += mev_buffer

        # Expand if revert rate is high (> 10%)
        if revert_rate > 0.10:
            slippage *= self.REVERT_ADJUST_FACTOR * (1 + revert_rate)

        # Urgency multiplier
        slippage *= urgency

        return max(self.MIN_SLIPPAGE_BPS, min(self.MAX_SLIPPAGE_BPS, slippage))

import asyncio
import aiohttp
import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from enum import Enum

logger = logging.getLogger("SmartOrderRouter")


class MEVProtection(Enum):
    NONE = "none"
    FLASHBOTS = "flashbots"
    MEV_BLOCKER = "mev_blocker"
    COWSWAP = "cowswap"


@dataclass
class QuoteResult:
    protocol: str
    input_token: str
    output_token: str
    input_amount: float
    output_amount: float
    price_impact_bps: float
    gas_estimate: int
    route_path: List[str]
    fee_bps: float
    data: str = ""  # encoded calldata for execution

    @property
    def effective_rate(self) -> float:
        return self.output_amount / self.input_amount

    def gas_adjusted_output(self, gas_price_gwei: float, eth_price_usd: float) -> float:
        gas_cost_usd = self.gas_estimate * gas_price_gwei * 1e-9 * eth_price_usd
        return self.output_amount - gas_cost_usd


@dataclass
class OrderConfig:
    input_token: str           # e.g. "USDC"
    output_token: str          # e.g. "ETH"
    input_amount_usd: float   # in USD terms
    max_splits: int = 3       # max number of routes to split across
    max_hops: int = 3         # max hops per route
    max_impact_bps: float = 100  # reject routes with > 1% impact
    mev_protection: MEVProtection = MEVProtection.FLASHBOTS
    urgency: float = 1.0
    chain_id: int = 1         # 1=Ethereum, 42161=Arbitrum, 10=Optimism


@dataclass
class ExecutionPlan:
    legs: List[Tuple[QuoteResult, float]]  # (quote, allocation_fraction)
    total_input: float
    total_output_estimate: float
    blended_impact_bps: float
    slippage_tolerance_bps: float
    mev_protection: MEVProtection
    estimated_gas: int

    def summary(self) -> str:
        lines = [f"ExecutionPlan: {len(self.legs)} legs"]
        for q, frac in self.legs:
            lines.append(
                f"  [{q.protocol}] {frac*100:.1f}% via {' → '.join(q.route_path)}"
                f" | impact={q.price_impact_bps:.1f}bps | gas={q.gas_estimate:,}"
            )
        lines.append(f"  Total output: {self.total_output_estimate:.4f}")
        lines.append(f"  Slippage tol: {self.slippage_tolerance_bps:.0f}bps")
        lines.append(f"  MEV protection: {self.mev_protection.value}")
        return "\n".join(lines)


class SmartOrderRouter:
    """
    Production-ready smart order router for AI trading agents.

    Aggregates quotes from multiple DEX aggregators, computes optimal
    split allocation, applies gas-adjusted scoring, and packages an
    execution plan with dynamic slippage and MEV protection config.

    Integrates with Purple Flea Trading API for on-chain execution.
    """

    # Aggregator API endpoints
    AGGREGATORS = {
        "1inch": "https://api.1inch.dev/swap/v6.0/{chain}/quote",
        "paraswap": "https://apiv5.paraswap.io/prices",
        "0x": "https://api.0x.org/swap/v1/quote",
        "kyberswap": "https://aggregator-api.kyberswap.com/{chain}/api/v1/routes",
    }

    MEV_RPCS = {
        MEVProtection.FLASHBOTS: "https://rpc.flashbots.net",
        MEVProtection.MEV_BLOCKER: "https://rpc.mevblocker.io",
        MEVProtection.COWSWAP: "https://api.cow.fi/mainnet/api/v1/quote",
        MEVProtection.NONE: "https://eth.llamarpc.com",
    }

    def __init__(
        self,
        api_keys: Optional[Dict[str, str]] = None,
        purple_flea_api_key: Optional[str] = None,
        eth_price_usd: float = 3200.0,
    ):
        self.api_keys = api_keys or {}
        self.pf_key = purple_flea_api_key
        self.eth_price_usd = eth_price_usd
        self.slippage_calc = DynamicSlippageCalculator()
        self._session: Optional[aiohttp.ClientSession] = None

    async def _session_get(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=5)
            )
        return self._session

    async def _fetch_quote_1inch(
        self, config: OrderConfig, amount: float
    ) -> Optional[QuoteResult]:
        try:
            session = await self._session_get()
            url = self.AGGREGATORS["1inch"].format(chain=config.chain_id)
            params = {
                "src": config.input_token,
                "dst": config.output_token,
                "amount": str(int(amount * 1e6)),  # assume USDC 6-decimals
                "includeGas": "true",
            }
            headers = {"Authorization": f"Bearer {self.api_keys.get('1inch', '')}"}
            async with session.get(url, params=params, headers=headers) as r:
                if r.status != 200:
                    return None
                data = await r.json()
                output = float(data["toAmount"]) / 1e18
                return QuoteResult(
                    protocol="1inch",
                    input_token=config.input_token,
                    output_token=config.output_token,
                    input_amount=amount,
                    output_amount=output,
                    price_impact_bps=float(data.get("priceImpact", 0)) * 10000,
                    gas_estimate=int(data.get("gas", 180000)),
                    route_path=data.get("route", [config.input_token, config.output_token]),
                    fee_bps=0,
                )
        except Exception as e:
            logger.warning(f"1inch quote failed: {e}")
            return None

    async def _fetch_all_quotes(
        self, config: OrderConfig, amount: float
    ) -> List[QuoteResult]:
        # Run all aggregator queries concurrently
        tasks = [self._fetch_quote_1inch(config, amount)]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        quotes = [r for r in results if isinstance(r, QuoteResult)]
        # Filter by max impact
        quotes = [q for q in quotes if q.price_impact_bps <= config.max_impact_bps]
        # Sort by gas-adjusted output descending
        gas_price_gwei = 20  # fetch from RPC in production
        quotes.sort(
            key=lambda q: q.gas_adjusted_output(gas_price_gwei, self.eth_price_usd),
            reverse=True
        )
        return quotes

    async def build_execution_plan(self, config: OrderConfig) -> ExecutionPlan:
        """
        Core routing method: fetches quotes, computes optimal split,
        calculates slippage, selects MEV protection.
        """
        logger.info(f"Building execution plan: {config.input_amount_usd} USD {config.input_token}→{config.output_token}")

        # Fetch quotes for full amount and for split fractions
        quotes = await self._fetch_all_quotes(config, config.input_amount_usd)

        if not quotes:
            raise RuntimeError("No valid quotes available from any aggregator")

        # Simple greedy: use top-k quotes with equal initial split
        top_k = quotes[:min(config.max_splits, len(quotes))]
        n = len(top_k)
        allocations = [config.input_amount_usd / n] * n

        # Compute blended impact
        blended_impact = sum(
            q.price_impact_bps * (a / config.input_amount_usd)
            for q, a in zip(top_k, allocations)
        )

        # Dynamic slippage
        slippage = self.slippage_calc.calculate(
            price_impact_bps=blended_impact,
            mev_protected=(config.mev_protection != MEVProtection.NONE),
            urgency=config.urgency,
        )

        # Estimate total output
        total_output = sum(
            q.output_amount * (a / config.input_amount_usd)
            for q, a in zip(top_k, allocations)
        )

        total_gas = sum(q.gas_estimate for q in top_k)

        plan = ExecutionPlan(
            legs=list(zip(top_k, [a / config.input_amount_usd for a in allocations])),
            total_input=config.input_amount_usd,
            total_output_estimate=total_output,
            blended_impact_bps=blended_impact,
            slippage_tolerance_bps=slippage,
            mev_protection=config.mev_protection,
            estimated_gas=total_gas,
        )

        logger.info(plan.summary())
        return plan

    async def close(self) -> None:
        if self._session:
            await self._session.close()


# Example usage
async def main():
    router = SmartOrderRouter(
        api_keys={"1inch": "YOUR_1INCH_API_KEY"},
        purple_flea_api_key="YOUR_PF_API_KEY",
    )
    plan = await router.build_execution_plan(OrderConfig(
        input_token="USDC",
        output_token="ETH",
        input_amount_usd=50000,
        max_splits=3,
        mev_protection=MEVProtection.FLASHBOTS,
    ))
    print(plan.summary())
    await router.close()

if __name__ == "__main__":
    asyncio.run(main())

import httpx

PF_TRADING_URL = "https://trading.purpleflea.com/api/v1"
PF_WALLET_URL = "https://wallet.purpleflea.com/api/v1"

async def execute_via_purple_flea(
    api_key: str,
    wallet_id: str,
    input_token: str,
    output_token: str,
    amount_usd: float,
    slippage_bps: int = 50,
):
    """Submit a trade order through Purple Flea's routing layer."""
    async with httpx.AsyncClient() as client:
        # Step 1: Get balance from wallet
        balance_resp = await client.get(
            f"{PF_WALLET_URL}/balance/{wallet_id}",
            headers={"X-API-Key": api_key},
        )
        balance_resp.raise_for_status()
        balance = balance_resp.json()["balances"]

        # Step 2: Get routing quote
        quote_resp = await client.post(
            f"{PF_TRADING_URL}/quote",
            headers={"X-API-Key": api_key},
            json={
                "from_token": input_token,
                "to_token": output_token,
                "amount_usd": amount_usd,
                "slippage_bps": slippage_bps,
                "mev_protection": "flashbots",
            },
        )
        quote_resp.raise_for_status()
        quote = quote_resp.json()

        print(f"Quote: {quote['estimated_output']} {output_token}")
        print(f"Route: {' → '.join(quote['route_path'])}")
        print(f"Impact: {quote['price_impact_bps']}bps")

        # Step 3: Execute
        exec_resp = await client.post(
            f"{PF_TRADING_URL}/execute",
            headers={"X-API-Key": api_key},
            json={
                "quote_id": quote["quote_id"],
                "wallet_id": wallet_id,
            },
        )
        exec_resp.raise_for_status()
        result = exec_resp.json()

        print(f"Executed: tx={result['tx_hash']}")
        print(f"Actual output: {result['actual_output']} {output_token}")
        print(f"Actual slippage: {result['actual_slippage_bps']}bps")
        return result