prediction-market-agent-tooling 0.65.5__py3-none-any.whl → 0.69.17.dev1149__py3-none-any.whl

prediction_market_agent_tooling/tools/contract_utils.py

@@ -0,0 +1,61 @@
+from web3 import Web3
+
+from prediction_market_agent_tooling.config import RPCConfig
+from prediction_market_agent_tooling.gtypes import ChecksumAddress
+from prediction_market_agent_tooling.tools.contract import contract_implements_function
+
+
+def is_erc20_contract(address: ChecksumAddress, web3: Web3 | None = None) -> bool:
+    """
+    Checks if the given address is an ERC20-compatible contract.
+
+    It estimates it by looking if it implements symbol, name, decimals, totalSupply, balanceOf.
+    """
+    web3 = web3 or RPCConfig().get_web3()
+    return (
+        contract_implements_function(address, "symbol", web3)
+        and contract_implements_function(address, "name", web3)
+        and contract_implements_function(address, "totalSupply", web3)
+        and contract_implements_function(
+            address, "balanceOf", web3, function_arg_types=["address"]
+        )
+    )
+
+
+def is_nft_contract(address: ChecksumAddress, web3: Web3 | None = None) -> bool:
+    """
+    Checks if the given address is an NFT-compatible contract (ERC721 or ERC1155).
+
+    For ERC721, checks for: ownerOf, balanceOf, transferFrom.
+    For ERC1155, checks for: balanceOf, safeTransferFrom.
+
+    Returns True if either ERC721 or ERC1155 interface is detected.
+    """
+    web3 = web3 or RPCConfig().get_web3()
+    is_erc721 = (
+        contract_implements_function(
+            address, "ownerOf", web3, function_arg_types=["uint256"]
+        )
+        and contract_implements_function(
+            address, "balanceOf", web3, function_arg_types=["address"]
+        )
+        and contract_implements_function(
+            address,
+            "transferFrom",
+            web3,
+            function_arg_types=["address", "address", "uint256"],
+        )
+    )
+    if is_erc721:
+        return True
+    is_erc1155 = contract_implements_function(
+        address, "balanceOf", web3, function_arg_types=["address", "uint256"]
+    ) and contract_implements_function(
+        address,
+        "safeTransferFrom",
+        web3,
+        function_arg_types=["address", "address", "uint256", "uint256", "bytes"],
+    )
+    if is_erc1155:
+        return True
+    return False

prediction_market_agent_tooling/tools/cow/cow_order.py

@@ -1,4 +1,5 @@
 import asyncio
+import typing as t
 from datetime import timedelta
 
 import httpx
@@ -8,11 +9,17 @@ from cowdao_cowpy.common.api.errors import UnexpectedResponseError
 from cowdao_cowpy.common.chains import Chain
 from cowdao_cowpy.common.config import SupportedChainId
 from cowdao_cowpy.common.constants import CowContractAddress
+from cowdao_cowpy.contracts.domain import domain
+from cowdao_cowpy.contracts.sign import SigningScheme, sign_order_cancellations
 from cowdao_cowpy.cow.swap import CompletedOrder, get_order_quote, swap_tokens
 from cowdao_cowpy.order_book.api import OrderBookApi
 from cowdao_cowpy.order_book.config import Envs, OrderBookAPIConfigFactory
 from cowdao_cowpy.order_book.generated.model import (
+    UID,
     Address,
+    EcdsaSignature,
+    EcdsaSigningScheme,
+    OrderCancellations,
     OrderMetaData,
     OrderQuoteRequest,
     OrderQuoteResponse,
@@ -23,28 +30,28 @@ from cowdao_cowpy.order_book.generated.model import (
     OrderStatus,
     TokenAmount,
 )
-from cowdao_cowpy.subgraph.client import BaseModel
-from tenacity import (
-    retry_if_not_exception_type,
-    stop_after_attempt,
-    wait_exponential,
-    wait_fixed,
-)
+from eth_account import Account
+from eth_account.signers.local import LocalAccount
+from eth_keys.datatypes import PrivateKey as eth_keys_PrivateKey
+from tenacity import stop_after_attempt, wait_exponential, wait_fixed
 from web3 import Web3
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.gtypes import ChecksumAddress, HexBytes, Wei
+from prediction_market_agent_tooling.gtypes import (
+    ChecksumAddress,
+    HexBytes,
+    HexStr,
+    PrivateKey,
+    Wei,
+)
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.omen.cow_contracts import (
     CowGPv2SettlementContract,
 )
-from prediction_market_agent_tooling.tools.contract import ContractERC20OnGnosisChain
-from prediction_market_agent_tooling.tools.utils import check_not_none, utcnow
-
-
-class MinimalisticToken(BaseModel):
-    sellToken: ChecksumAddress
-    buyToken: ChecksumAddress
+from prediction_market_agent_tooling.tools.contract import ContractERC20BaseClass
+from prediction_market_agent_tooling.tools.cow.models import MinimalisticTrade, Order
+from prediction_market_agent_tooling.tools.cow.semaphore import postgres_rate_limited
+from prediction_market_agent_tooling.tools.utils import utcnow
 
 
 class OrderStatusError(Exception):
@@ -96,7 +103,6 @@ def get_sell_token_amount(
 @tenacity.retry(
     stop=stop_after_attempt(4),
     wait=wait_exponential(min=4, max=10),
-    retry=retry_if_not_exception_type(NoLiquidityAvailableOnCowException),
 )
 def get_quote(
     amount_wei: Wei,
@@ -157,10 +163,41 @@ def get_buy_token_amount_else_raise(
     return Wei(order_quote.quote.buyAmount.root)
 
 
+def handle_allowance(
+    api_keys: APIKeys,
+    sell_token: ChecksumAddress,
+    amount_to_check_wei: Wei,
+    amount_to_set_wei: Wei | None = None,
+    for_address: ChecksumAddress | None = None,
+    web3: Web3 | None = None,
+) -> None:
+    # Approve the CoW Swap Vault Relayer to get the sell token only if allowance not sufficient.
+    for_address = for_address or Web3.to_checksum_address(
+        CowContractAddress.VAULT_RELAYER.value
+    )
+    current_allowance = ContractERC20BaseClass(address=sell_token).allowance(
+        owner=api_keys.bet_from_address,
+        for_address=for_address,
+        web3=web3,
+    )
+    if current_allowance < amount_to_check_wei:
+        amount_to_set_wei = amount_to_set_wei or amount_to_check_wei
+        ContractERC20BaseClass(address=sell_token).approve(
+            api_keys,
+            for_address=for_address,
+            amount_wei=amount_to_set_wei,
+            web3=web3,
+        )
+
+
+@postgres_rate_limited(
+    api_keys=APIKeys(), rate_id="swap_tokens_waiting", interval_seconds=5.0
+)
 @tenacity.retry(
+    reraise=True,
     stop=stop_after_attempt(3),
     wait=wait_fixed(1),
-    retry=tenacity.retry_if_not_exception_type((TimeoutError, OrderStatusError)),
+    retry=tenacity.retry_if_not_exception_type((TimeoutError)),
     after=lambda x: logger.debug(f"swap_tokens_waiting failed, {x.attempt_number=}."),
 )
 def swap_tokens_waiting(
@@ -171,40 +208,37 @@ def swap_tokens_waiting(
     chain: Chain = Chain.GNOSIS,
     env: Envs = "prod",
     web3: Web3 | None = None,
-) -> OrderMetaData:
-    # Approve the CoW Swap Vault Relayer to get the sell token only if allowance not sufficient.
-    for_address = Web3.to_checksum_address(CowContractAddress.VAULT_RELAYER.value)
-    current_allowance = ContractERC20OnGnosisChain(address=sell_token).allowance(
-        owner=api_keys.bet_from_address,
-        for_address=for_address,
-        web3=web3,
-    )
-    if current_allowance < amount_wei:
-        ContractERC20OnGnosisChain(address=sell_token).approve(
-            api_keys,
-            for_address=for_address,
-            amount_wei=amount_wei,
-            web3=web3,
-        )
-
+    wait_order_complete: bool = True,
+    timeout: timedelta = timedelta(seconds=120),
+    slippage_tolerance: float = 0.01,
+) -> tuple[OrderMetaData | None, CompletedOrder]:
     # CoW library uses async, so we need to wrap the call in asyncio.run for us to use it.
     return asyncio.run(
         swap_tokens_waiting_async(
-            amount_wei, sell_token, buy_token, api_keys, chain, env
+            amount_wei,
+            sell_token,
+            buy_token,
+            api_keys,
+            chain,
+            env,
+            timeout=timeout,
+            web3=web3,
+            wait_order_complete=wait_order_complete,
+            slippage_tolerance=slippage_tolerance,
         )
     )
 
 
-async def swap_tokens_waiting_async(
+async def place_swap_order(
+    api_keys: APIKeys,
     amount_wei: Wei,
     sell_token: ChecksumAddress,
     buy_token: ChecksumAddress,
-    api_keys: APIKeys,
     chain: Chain,
     env: Envs,
-    timeout: timedelta = timedelta(seconds=120),
     slippage_tolerance: float = 0.01,
-) -> OrderMetaData:
+    valid_to: int | None = None,
+) -> CompletedOrder:
     account = api_keys.get_account()
     safe_address = api_keys.safe_address_checksum
 
@@ -217,6 +251,7 @@ async def swap_tokens_waiting_async(
         chain=chain,
         env=env,
         slippage_tolerance=slippage_tolerance,
+        valid_to=valid_to,
     )
     logger.info(f"Order created: {order}")
 
@@ -224,6 +259,12 @@ async def swap_tokens_waiting_async(
         logger.info(f"Safe is used. Signing the order after its creation.")
         await sign_safe_cow_swap(api_keys, order, chain, env)
 
+    return order
+
+
+async def wait_for_order_completion(
+    order: CompletedOrder, timeout: timedelta = timedelta(seconds=120)
+) -> OrderMetaData:
     start_time = utcnow()
 
     while True:
@@ -253,6 +294,41 @@ async def swap_tokens_waiting_async(
         await asyncio.sleep(3.14)
 
 
+async def swap_tokens_waiting_async(
+    amount_wei: Wei,
+    sell_token: ChecksumAddress,
+    buy_token: ChecksumAddress,
+    api_keys: APIKeys,
+    chain: Chain,
+    env: Envs,
+    timeout: timedelta = timedelta(seconds=120),
+    slippage_tolerance: float = 0.01,
+    web3: Web3 | None = None,
+    wait_order_complete: bool = True,
+) -> tuple[OrderMetaData | None, CompletedOrder]:
+    handle_allowance(
+        api_keys=api_keys,
+        sell_token=sell_token,
+        amount_to_check_wei=amount_wei,
+        web3=web3,
+    )
+    valid_to = (utcnow() + timeout).timestamp()
+    order = await place_swap_order(
+        api_keys=api_keys,
+        amount_wei=amount_wei,
+        sell_token=sell_token,
+        buy_token=buy_token,
+        chain=chain,
+        env=env,
+        slippage_tolerance=slippage_tolerance,
+        valid_to=int(valid_to),
+    )
+    if wait_order_complete:
+        order_metadata = await wait_for_order_completion(order=order, timeout=timeout)
+        return order_metadata, order
+    return None, order
+
+
 @tenacity.retry(
     stop=tenacity.stop_after_attempt(3),
     wait=tenacity.wait_fixed(1),
@@ -266,11 +342,8 @@ async def sign_safe_cow_swap(
 ) -> None:
     order_book_api = get_order_book_api(env, chain)
     posted_order = await order_book_api.get_order_by_uid(order.uid)
-    CowGPv2SettlementContract(
-        address=Web3.to_checksum_address(
-            check_not_none(posted_order.settlementContract).root
-        )
-    ).setPreSignature(
+
+    CowGPv2SettlementContract().setPreSignature(
         api_keys,
         HexBytes(posted_order.uid.root),
         True,
@@ -284,11 +357,111 @@ async def sign_safe_cow_swap(
 )
 def get_trades_by_owner(
     owner: ChecksumAddress,
-) -> list[MinimalisticToken]:
+) -> list[MinimalisticTrade]:
     # Using this until cowpy gets fixed (https://github.com/cowdao-grants/cow-py/issues/35)
     response = httpx.get(
         f"https://api.cow.fi/xdai/api/v1/trades",
         params={"owner": owner},
     )
     response.raise_for_status()
-    return [MinimalisticToken.model_validate(i) for i in response.json()]
+    return [MinimalisticTrade.model_validate(i) for i in response.json()]
+
+
+@tenacity.retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(1),
+    after=lambda x: logger.debug(
+        f"get_trades_by_order_uid failed, {x.attempt_number=}."
+    ),
+)
+def get_trades_by_order_uid(
+    order_uid: HexBytes,
+) -> list[MinimalisticTrade]:
+    # Using this until cowpy gets fixed (https://github.com/cowdao-grants/cow-py/issues/35)
+    response = httpx.get(
+        f"https://api.cow.fi/xdai/api/v1/trades",
+        params={"orderUid": order_uid.to_0x_hex()},
+    )
+    response.raise_for_status()
+    return [MinimalisticTrade.model_validate(i) for i in response.json()]
+
+
+@tenacity.retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(1),
+    after=lambda x: logger.debug(f"get_orders_by_owner failed, {x.attempt_number=}."),
+)
+def get_orders_by_owner(
+    owner: ChecksumAddress,
+) -> list[Order]:
+    """Retrieves all orders with pagination."""
+    items = paginate_endpoint(f"https://api.cow.fi/xdai/api/v1/account/{owner}/orders")
+    return [Order.model_validate(i) for i in items]
+
+
+@tenacity.retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(1),
+    after=lambda x: logger.debug(f"get_order_by_uid failed, {x.attempt_number=}."),
+)
+async def get_order_by_uid(
+    uid: HexBytes,
+) -> Order:
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            f"https://api.cow.fi/xdai/api/v1/orders/{uid.to_0x_hex()}"
+        )
+        response.raise_for_status()
+        return Order.model_validate(response.json())
+
+
+@tenacity.retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(1),
+    after=lambda x: logger.debug(f"paginate_endpoint failed, {x.attempt_number=}."),
+)
+def paginate_endpoint(url: str, limit: int = 1000) -> t.Any:
+    offset = 0
+    results = []
+
+    while True:
+        response = httpx.get(url, params={"offset": offset, "limit": limit})
+        response.raise_for_status()
+        items = response.json()
+        if not items:
+            break
+
+        results.extend(items)
+        offset += limit
+
+    return results
+
+
+async def cancel_order(
+    order_uids: t.Sequence[str],
+    api_keys: APIKeys,
+    chain: Chain = Chain.GNOSIS,
+    env: Envs = "prod",
+) -> None:
+    """Cancels orders that were created by the user corresponding to the api_keys.."""
+    order_domain = domain(
+        chain=chain, verifying_contract=CowContractAddress.SETTLEMENT_CONTRACT.value
+    )
+    owner = build_account(api_keys.bet_from_private_key)
+    order_signature = sign_order_cancellations(
+        order_domain, list(order_uids), owner, SigningScheme.EIP712
+    )
+    oc = OrderCancellations(
+        signature=EcdsaSignature(root=order_signature.data),
+        orderUids=[UID(root=order_uid) for order_uid in order_uids],
+        signingScheme=EcdsaSigningScheme.eip712,
+    )
+    order_book_api = get_order_book_api(env, chain)
+    await order_book_api.delete_order(orders_cancelation=oc)
+
+
+def build_account(private_key: PrivateKey) -> LocalAccount:
+    return LocalAccount(
+        key=eth_keys_PrivateKey(HexBytes(HexStr(private_key.get_secret_value()))),
+        account=Account.from_key(private_key.get_secret_value()),
+    )

prediction_market_agent_tooling/tools/cow/models.py

@@ -0,0 +1,122 @@
+from enum import Enum
+from typing import Optional, TypeAlias
+
+from pydantic import BaseModel, ConfigDict
+from sqlmodel import Field, SQLModel
+
+from prediction_market_agent_tooling.gtypes import (
+    HexBytes,
+    VerifiedChecksumAddress,
+    VerifiedChecksumAddressOrNone,
+    Wei,
+)
+from prediction_market_agent_tooling.tools.datetime_utc import DatetimeUTC
+from prediction_market_agent_tooling.tools.utils import utcnow
+
+
+class MinimalisticTrade(BaseModel):
+    sellToken: VerifiedChecksumAddress
+    buyToken: VerifiedChecksumAddress
+    orderUid: HexBytes
+    txHash: HexBytes
+
+
+class EthFlowData(BaseModel):
+    refundTxHash: Optional[HexBytes]
+    userValidTo: int
+
+
+class PlacementError(str, Enum):
+    QuoteNotFound = "QuoteNotFound"
+    ValidToTooFarInFuture = "ValidToTooFarInFuture"
+    PreValidationError = "PreValidationError"
+
+
+class OnchainOrderData(BaseModel):
+    sender: VerifiedChecksumAddress
+    placementError: Optional[PlacementError]
+
+
+class OrderKind(str, Enum):
+    buy = "buy"
+    sell = "sell"
+
+
+class SellTokenBalance(str, Enum):
+    external = "external"
+    internal = "internal"
+    erc20 = "erc20"
+
+
+class BuyTokenBalance(str, Enum):
+    internal = "internal"
+    erc20 = "erc20"
+
+
+class SigningScheme(str, Enum):
+    eip712 = "eip712"
+    ethsign = "ethsign"
+    presign = "presign"
+    eip1271 = "eip1271"
+
+
+class OrderClass(str, Enum):
+    limit = "limit"
+    liquidity = "liquidity"
+    market = "market"
+
+
+class OrderStatus(str, Enum):
+    presignaturePending = "presignaturePending"
+    open = "open"
+    fulfilled = "fulfilled"
+    cancelled = "cancelled"
+    expired = "expired"
+
+
+CowOrderUID: TypeAlias = HexBytes
+
+
+class Order(BaseModel):
+    model_config = ConfigDict(validate_by_name=True, validate_by_alias=True)
+
+    uid: CowOrderUID
+    quoteId: int | None = None
+    validTo: int
+    sellAmount: Wei
+    sellToken: VerifiedChecksumAddress
+    buyAmount: Wei
+    buyToken: VerifiedChecksumAddress
+    receiver: VerifiedChecksumAddressOrNone
+    feeAmount: Wei
+    creationDate: DatetimeUTC
+    kind: OrderKind
+    partiallyFillable: bool
+    sellTokenBalance: SellTokenBalance | None
+    buyTokenBalance: BuyTokenBalance | None
+    signingScheme: SigningScheme
+    signature: HexBytes
+    from_: VerifiedChecksumAddressOrNone = Field(None, alias="from")
+    appData: str
+    fullAppData: str | None
+    appDataHash: HexBytes | None = None
+    class_: str = Field(None, alias="class")
+    owner: VerifiedChecksumAddress
+    executedSellAmount: Wei
+    executedSellAmountBeforeFees: Wei
+    executedBuyAmount: Wei
+    executedFeeAmount: Wei | None
+    invalidated: bool
+    status: str
+    isLiquidityOrder: bool | None
+    ethflowData: EthFlowData | None = None
+    onchainUser: VerifiedChecksumAddressOrNone = None
+    executedFee: Wei
+    executedFeeToken: VerifiedChecksumAddressOrNone
+
+
+class RateLimit(SQLModel, table=True):
+    __tablename__ = "rate_limit"
+    __table_args__ = {"extend_existing": True}
+    id: str = Field(primary_key=True)
+    last_called_at: DatetimeUTC = Field(default_factory=utcnow)

prediction_market_agent_tooling/tools/cow/semaphore.py

@@ -0,0 +1,104 @@
+import time
+from datetime import timedelta
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar, cast
+
+from sqlalchemy.exc import OperationalError
+from sqlmodel import Session, select
+
+from prediction_market_agent_tooling.config import APIKeys
+from prediction_market_agent_tooling.tools.cow.models import RateLimit
+from prediction_market_agent_tooling.tools.datetime_utc import DatetimeUTC
+from prediction_market_agent_tooling.tools.db.db_manager import DBManager
+from prediction_market_agent_tooling.tools.utils import utcnow
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+FALLBACK_SQL_ENGINE = "sqlite:///rate_limit.db"
+
+
+def postgres_rate_limited(
+    api_keys: APIKeys,
+    rate_id: str,
+    interval_seconds: float,
+    shared_db: bool = False,
+) -> Callable[[F], F]:
+    """rate_id is used to distinguish between different rate limits for different functions"""
+    limiter = RateLimiter(id=rate_id, interval_seconds=interval_seconds)
+
+    def decorator(func: F) -> F:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            sqlalchemy_db_url = (
+                api_keys.sqlalchemy_db_url.get_secret_value()
+                if shared_db
+                else FALLBACK_SQL_ENGINE
+            )
+
+            db_manager = DBManager(sqlalchemy_db_url)
+            db_manager.create_tables([RateLimit])
+
+            with db_manager.get_session() as session:
+                limiter.enforce(session)
+            return func(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    return decorator
+
+
+class RateLimiter:
+    def __init__(self, id: str, interval_seconds: float = 1.0) -> None:
+        self.id = id
+        self.interval = timedelta(seconds=interval_seconds)
+
+    def enforce(self, session: Session, timeout_seconds: float = 30.0) -> None:
+        """
+        Enforces the rate limit inside a transaction.
+        Blocks until allowed or timeout is reached.
+
+        Args:
+            session: The database session to use
+            timeout_seconds: Maximum time in seconds to wait before giving up
+
+        Raises:
+            TimeoutError: If the rate limit cannot be acquired within the timeout period
+        """
+        start_time = time.monotonic()
+
+        while True:
+            try:
+                with session.begin():
+                    stmt = (
+                        select(RateLimit)
+                        .where(RateLimit.id == self.id)
+                        .with_for_update()
+                    )
+                    result: Optional[RateLimit] = session.exec(stmt).first()
+
+                    now = utcnow()
+
+                    if result is None:
+                        # First time this limiter is used
+                        session.add(RateLimit(id=self.id))
+                        return
+
+                    last_called_aware = DatetimeUTC.from_datetime(result.last_called_at)
+                    elapsed = now - last_called_aware
+                    if elapsed >= self.interval:
+                        result.last_called_at = now
+                        session.add(result)
+                        return
+
+                    # Not enough time passed, sleep and retry
+                    to_sleep = (self.interval - elapsed).total_seconds()
+                    time.sleep(to_sleep)
+            except OperationalError:
+                # Backoff if DB is under contention
+                elapsed_time = time.monotonic() - start_time
+                if elapsed_time > timeout_seconds:
+                    raise TimeoutError(
+                        f"Could not acquire rate limit '{self.id}' "
+                        f"after {elapsed_time:.1f} seconds due to database contention"
+                    )
+                time.sleep(0.5)

prediction_market_agent_tooling/tools/datetime_utc.py

@@ -34,8 +34,20 @@ class DatetimeUTC(datetime):
     def __get_pydantic_core_schema__(
         cls, source_type: t.Any, handler: GetCoreSchemaHandler
     ) -> CoreSchema:
-        dt_schema = handler(datetime)
-        return core_schema.no_info_after_validator_function(cls._validate, dt_schema)
+        # Use union schema to handle int, str, and datetime inputs directly
+        return core_schema.union_schema(
+            [
+                core_schema.no_info_after_validator_function(
+                    cls._validate, core_schema.int_schema()
+                ),
+                core_schema.no_info_after_validator_function(
+                    cls._validate, core_schema.str_schema()
+                ),
+                core_schema.no_info_after_validator_function(
+                    cls._validate, handler(datetime)
+                ),
+            ]
+        )
 
     @staticmethod
     def from_datetime(dt: datetime) -> "DatetimeUTC":