prediction-market-agent-tooling 0.61.2.dev479__py3-none-any.whl → 0.62.0__py3-none-any.whl

prediction_market_agent_tooling/tools/balances.py

@@ -1,32 +1,35 @@
 from pydantic import BaseModel
 from tenacity import retry, stop_after_attempt, wait_fixed
 from web3 import Web3
-from web3.types import Wei
 
-from prediction_market_agent_tooling.gtypes import ChecksumAddress, xDai
+from prediction_market_agent_tooling.gtypes import (
+    ChecksumAddress,
+    CollateralToken,
+    xDai,
+    xDaiWei,
+)
 from prediction_market_agent_tooling.markets.omen.omen_contracts import (
     WrappedxDaiContract,
     sDaiContract,
 )
-from prediction_market_agent_tooling.tools.web3_utils import wei_to_xdai
 
 
 class Balances(BaseModel):
     xdai: xDai
-    wxdai: xDai
-    sdai: xDai
+    wxdai: CollateralToken
+    sdai: CollateralToken
 
     @property
-    def total(self) -> xDai:
-        return xDai(self.xdai + self.wxdai + self.sdai)
+    def total(self) -> CollateralToken:
+        return self.xdai.as_token + self.wxdai + self.sdai
 
 
 @retry(stop=stop_after_attempt(3), wait=wait_fixed(1))
 def get_balances(address: ChecksumAddress, web3: Web3 | None = None) -> Balances:
     if not web3:
         web3 = WrappedxDaiContract().get_web3()
-    xdai_balance = Wei(web3.eth.get_balance(address))
-    xdai = wei_to_xdai(xdai_balance)
-    wxdai = wei_to_xdai(WrappedxDaiContract().balanceOf(address, web3=web3))
-    sdai = wei_to_xdai(sDaiContract().balanceOf(address, web3=web3))
+    xdai_balance = xDaiWei(web3.eth.get_balance(address))
+    xdai = xdai_balance.as_xdai
+    wxdai = WrappedxDaiContract().balanceOf(address, web3=web3).as_token
+    sdai = sDaiContract().balanceOf(address, web3=web3).as_token
     return Balances(xdai=xdai, wxdai=wxdai, sdai=sdai)

prediction_market_agent_tooling/tools/betting_strategies/kelly_criterion.py

@@ -1,3 +1,4 @@
+from prediction_market_agent_tooling.gtypes import CollateralToken, OutcomeToken
 from prediction_market_agent_tooling.markets.market_fees import MarketFees
 from prediction_market_agent_tooling.tools.betting_strategies.utils import SimpleBet
 
@@ -8,7 +9,7 @@ def check_is_valid_probability(probability: float) -> None:
 
 
 def get_kelly_bet_simplified(
-    max_bet: float,
+    max_bet: CollateralToken,
     market_p_yes: float,
     estimated_p_yes: float,
     confidence: float,
@@ -51,17 +52,17 @@ def get_kelly_bet_simplified(
     kelly_fraction = edge / odds
 
     # Ensure bet size is non-negative does not exceed the wallet balance
-    bet_size = min(kelly_fraction * max_bet, max_bet)
+    bet_size = CollateralToken(min(kelly_fraction * max_bet.value, max_bet.value))
 
     return SimpleBet(direction=bet_direction, size=bet_size)
 
 
 def get_kelly_bet_full(
-    yes_outcome_pool_size: float,
-    no_outcome_pool_size: float,
+    yes_outcome_pool_size: OutcomeToken,
+    no_outcome_pool_size: OutcomeToken,
     estimated_p_yes: float,
     confidence: float,
-    max_bet: float,
+    max_bet: CollateralToken,
     fees: MarketFees,
 ) -> SimpleBet:
     """
@@ -97,13 +98,13 @@ def get_kelly_bet_full(
     check_is_valid_probability(confidence)
 
     if max_bet == 0:
-        return SimpleBet(direction=True, size=0)
+        return SimpleBet(direction=True, size=CollateralToken(0))
 
-    x = yes_outcome_pool_size
-    y = no_outcome_pool_size
+    x = yes_outcome_pool_size.value
+    y = no_outcome_pool_size.value
     p = estimated_p_yes
     c = confidence
-    b = max_bet
+    b = max_bet.value
     f = 1 - fee
 
     if x == y:
@@ -144,5 +145,6 @@ def get_kelly_bet_full(
 
     # Clip the bet size to max_bet to account for rounding errors.
     return SimpleBet(
-        direction=kelly_bet_amount > 0, size=min(max_bet, abs(kelly_bet_amount))
+        direction=kelly_bet_amount > 0,
+        size=CollateralToken(min(max_bet.value, abs(kelly_bet_amount))),
     )

prediction_market_agent_tooling/tools/betting_strategies/market_moving.py

@@ -2,19 +2,22 @@ from functools import reduce
 
 import numpy as np
 
-from prediction_market_agent_tooling.gtypes import Probability, Wei, xDai
+from prediction_market_agent_tooling.gtypes import (
+    CollateralToken,
+    OutcomeToken,
+    Probability,
+)
 from prediction_market_agent_tooling.markets.omen.omen import (
     MarketFees,
     OmenAgentMarket,
 )
 from prediction_market_agent_tooling.tools.betting_strategies.utils import SimpleBet
 from prediction_market_agent_tooling.tools.utils import check_not_none
-from prediction_market_agent_tooling.tools.web3_utils import wei_to_xdai, xdai_to_wei
 
 
 def get_market_moving_bet(
-    yes_outcome_pool_size: float,
-    no_outcome_pool_size: float,
+    yes_outcome_pool_size: OutcomeToken,
+    no_outcome_pool_size: OutcomeToken,
     market_p_yes: float,
     target_p_yes: float,
     fees: MarketFees,
@@ -42,20 +45,21 @@ def get_market_moving_bet(
     fixed_product = yes_outcome_pool_size * no_outcome_pool_size
     bet_direction: bool = target_p_yes > market_p_yes
 
-    min_bet_amount = 0.0
-    max_bet_amount = 100 * (
+    min_bet_amount = CollateralToken(0.0)
+    max_bet_amount = (
         yes_outcome_pool_size + no_outcome_pool_size
-    )  # TODO set a better upper bound
+    ).as_token * 100  # TODO set a better upper bound
 
     # Binary search for the optimal bet amount
     for _ in range(max_iters):
         bet_amount = (min_bet_amount + max_bet_amount) / 2
-        amounts_diff = fees.get_bet_size_after_fees(bet_amount)
+        amounts_diff = fees.get_after_fees(bet_amount)
+        amounts_diff_as_ot = OutcomeToken.from_token(amounts_diff)
 
         # Initial new amounts are old amounts + equal new amounts for each outcome
-        yes_outcome_new_pool_size = yes_outcome_pool_size + amounts_diff
-        no_outcome_new_pool_size = no_outcome_pool_size + amounts_diff
-        new_amounts = {
+        yes_outcome_new_pool_size = yes_outcome_pool_size + amounts_diff_as_ot
+        no_outcome_new_pool_size = no_outcome_pool_size + amounts_diff_as_ot
+        new_amounts: dict[bool, OutcomeToken] = {
             True: yes_outcome_new_pool_size,
             False: no_outcome_new_pool_size,
         }
@@ -63,15 +67,20 @@ def get_market_moving_bet(
         # Now give away tokens at `bet_outcome_index` to restore invariant
         new_product = yes_outcome_new_pool_size * no_outcome_new_pool_size
         dx = (new_product - fixed_product) / new_amounts[not bet_direction]
-        new_amounts[bet_direction] -= dx
+        new_amounts[bet_direction] -= OutcomeToken(dx)
 
         # Check that the invariant is restored
         assert np.isclose(
-            reduce(lambda x, y: x * y, list(new_amounts.values()), 1.0),
+            reduce(lambda x, y: x * y.value, list(new_amounts.values()), 1.0),
             float(fixed_product),
         )
 
-        new_p_yes = Probability(new_amounts[False] / sum(list(new_amounts.values())))
+        new_p_yes = Probability(
+            (
+                new_amounts[False]
+                / sum(list(new_amounts.values()), start=OutcomeToken(0))
+            )
+        )
         if abs(target_p_yes - new_p_yes) < 1e-6:
             break
         elif new_p_yes > target_p_yes:
@@ -97,33 +106,31 @@ def _sanity_check_omen_market_moving_bet(
     using the adjusted outcome pool sizes to calculate the new p_yes.
     """
     buy_amount_ = market.get_contract().calcBuyAmount(
-        investment_amount=xdai_to_wei(xDai(bet_to_check.size)),
+        investment_amount=bet_to_check.size.as_wei,
         outcome_index=market.get_outcome_index(
             market.get_outcome_str_from_bool(bet_to_check.direction)
         ),
     )
-    buy_amount = float(wei_to_xdai(Wei(buy_amount_)))
+    buy_amount = buy_amount_.as_outcome_token
 
     outcome_token_pool = check_not_none(market.outcome_token_pool)
     yes_outcome_pool_size = outcome_token_pool[market.get_outcome_str_from_bool(True)]
     no_outcome_pool_size = outcome_token_pool[market.get_outcome_str_from_bool(False)]
-    market_const = yes_outcome_pool_size * no_outcome_pool_size
+    market_const = yes_outcome_pool_size.value * no_outcome_pool_size.value
 
-    bet_to_check_size_after_fees = market.fees.get_bet_size_after_fees(
-        bet_to_check.size
-    )
+    bet_to_check_size_after_fees = market.fees.get_after_fees(bet_to_check.size).value
 
     # When you buy 'yes' tokens, you add your bet size to the both pools, then
     # subtract `buy_amount` from the 'yes' pool. And vice versa for 'no' tokens.
     new_yes_outcome_pool_size = (
-        yes_outcome_pool_size
+        yes_outcome_pool_size.value
         + bet_to_check_size_after_fees
-        - float(bet_to_check.direction) * buy_amount
+        - float(bet_to_check.direction) * buy_amount.value
     )
     new_no_outcome_pool_size = (
-        no_outcome_pool_size
+        no_outcome_pool_size.value
         + bet_to_check_size_after_fees
-        - float(not bet_to_check.direction) * buy_amount
+        - float(not bet_to_check.direction) * buy_amount.value
     )
     new_market_const = new_yes_outcome_pool_size * new_no_outcome_pool_size
     # Check the invariant is restored

prediction_market_agent_tooling/tools/betting_strategies/utils.py

@@ -1,6 +1,8 @@
 from pydantic import BaseModel
 
+from prediction_market_agent_tooling.gtypes import CollateralToken
+
 
 class SimpleBet(BaseModel):
     direction: bool
-    size: float
+    size: CollateralToken

prediction_market_agent_tooling/tools/contract.py

@@ -14,6 +14,7 @@ from prediction_market_agent_tooling.gtypes import (
     ABI,
     ChainID,
     ChecksumAddress,
+    CollateralToken,
     Nonce,
     TxParams,
     TxReceipt,
@@ -153,7 +154,7 @@ class ContractBaseClass(BaseModel):
             api_keys=api_keys,
             function_name=function_name,
             function_params=function_params,
-            tx_params={"value": amount_wei, **(tx_params or {})},
+            tx_params={"value": amount_wei.value, **(tx_params or {})},
             timeout=timeout,
             web3=web3,
         )
@@ -243,12 +244,13 @@ class ContractERC20BaseClass(ContractBaseClass):
         )
 
     def balanceOf(self, for_address: ChecksumAddress, web3: Web3 | None = None) -> Wei:
-        balance: Wei = self.call("balanceOf", [for_address], web3=web3)
+        balance = Wei(self.call("balanceOf", [for_address], web3=web3))
         return balance
 
-    def get_in_shares(self, amount: Wei, web3: Web3 | None = None) -> Wei:
-        # ERC-20 just holds the token, so the exact amount we send there, is the amount of shares we have there.
-        return amount
+    def balance_of_in_tokens(
+        self, for_address: ChecksumAddress, web3: Web3 | None = None
+    ) -> CollateralToken:
+        return self.balanceOf(for_address, web3=web3).as_token
 
 
 class ContractDepositableWrapperERC20BaseClass(ContractERC20BaseClass):
@@ -355,11 +357,11 @@ class ContractERC4626BaseClass(ContractERC20BaseClass):
         return self.withdraw(api_keys=api_keys, assets_wei=assets, web3=web3)
 
     def convertToShares(self, assets: Wei, web3: Web3 | None = None) -> Wei:
-        shares: Wei = self.call("convertToShares", [assets], web3=web3)
+        shares = Wei(self.call("convertToShares", [assets], web3=web3))
         return shares
 
     def convertToAssets(self, shares: Wei, web3: Web3 | None = None) -> Wei:
-        assets: Wei = self.call("convertToAssets", [shares], web3=web3)
+        assets = Wei(self.call("convertToAssets", [shares], web3=web3))
         return assets
 
     def get_asset_token_contract(
@@ -434,8 +436,8 @@ class ContractERC721BaseClass(ContractBaseClass):
             web3=web3,
         )
 
-    def balanceOf(self, owner: ChecksumAddress, web3: Web3 | None = None) -> int:
-        balance: int = self.call("balanceOf", [owner], web3=web3)
+    def balanceOf(self, owner: ChecksumAddress, web3: Web3 | None = None) -> Wei:
+        balance = Wei(self.call("balanceOf", [owner], web3=web3))
         return balance
 
     def owner_of(self, token_id: int, web3: Web3 | None = None) -> ChecksumAddress:
@@ -604,12 +606,14 @@ def contract_implements_function(
 
 
 def init_collateral_token_contract(
-    address: ChecksumAddress, web3: Web3
+    address: ChecksumAddress, web3: Web3 | None
 ) -> ContractERC20BaseClass:
     """
     Checks if the given contract is Depositable ERC-20, ERC-20 or ERC-4626 and returns the appropriate class instance.
     Throws an error if the contract is neither of them.
     """
+    web3 = web3 or RPCConfig().get_web3()
+
     if contract_implements_function(address, "asset", web3=web3):
         return ContractERC4626BaseClass(address=address)
 

prediction_market_agent_tooling/tools/cow/cow_manager.py

@@ -19,10 +19,9 @@ from web3 import Web3
 from web3.constants import ADDRESS_ZERO
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.gtypes import ChecksumAddress, Wei, xDai
+from prediction_market_agent_tooling.gtypes import ChecksumAddress, CollateralToken, Wei
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.cow.cow_order import swap_tokens_waiting
-from prediction_market_agent_tooling.tools.web3_utils import xdai_to_wei
 
 COW_ENV: Envs = "prod"
 
@@ -94,14 +93,14 @@ class CowManager:
 
     @staticmethod
     def swap(
-        amount: xDai,
+        amount: CollateralToken,
         sell_token: ChecksumAddress,
         buy_token: ChecksumAddress,
         api_keys: APIKeys,
         web3: Web3 | None = None,
     ) -> OrderMetaData:
         order_metadata = swap_tokens_waiting(
-            amount_wei=xdai_to_wei(amount),
+            amount_wei=amount.as_wei,
             sell_token=sell_token,
             buy_token=buy_token,
             api_keys=api_keys,

prediction_market_agent_tooling/tools/cow/cow_order.py

@@ -14,6 +14,8 @@ from cowdao_cowpy.order_book.generated.model import (
     OrderMetaData,
     OrderQuoteRequest,
     OrderQuoteSide1,
+    OrderQuoteSide3,
+    OrderQuoteSideKindBuy,
     OrderQuoteSideKindSell,
     OrderStatus,
     TokenAmount,
@@ -21,27 +23,63 @@ from cowdao_cowpy.order_book.generated.model import (
 from eth_account.signers.local import LocalAccount
 from eth_typing.evm import ChecksumAddress
 from web3 import Web3
-from web3.types import Wei
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.gtypes import ChecksumAddress, Wei, wei_type
+from prediction_market_agent_tooling.gtypes import ChecksumAddress, Wei
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.contract import ContractERC20OnGnosisChain
 from prediction_market_agent_tooling.tools.utils import utcnow
 
 
+class OrderStatusError(Exception):
+    pass
+
+
 def get_order_book_api(env: Envs, chain: Chain) -> OrderBookApi:
     chain_id = SupportedChainId(chain.value[0])
     return OrderBookApi(OrderBookAPIConfigFactory.get_config(env, chain_id))
 
 
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_fixed(1),
+    after=lambda x: logger.debug(f"get_sell_token_amount failed, {x.attempt_number=}."),
+)
+def get_sell_token_amount(
+    buy_amount: Wei,
+    sell_token: ChecksumAddress,
+    buy_token: ChecksumAddress,
+    chain: Chain = Chain.GNOSIS,
+    env: Envs = "prod",
+) -> Wei:
+    """
+    Calculate how much of the sell_token is needed to obtain a specified amount of buy_token.
+    """
+    order_book_api = get_order_book_api(env, chain)
+    order_quote_request = OrderQuoteRequest(
+        sellToken=Address(sell_token),
+        buyToken=Address(buy_token),
+        from_=Address(
+            "0x1234567890abcdef1234567890abcdef12345678"
+        ),  # Just random address, doesn't matter.
+    )
+    order_side = OrderQuoteSide3(
+        kind=OrderQuoteSideKindBuy.buy,
+        buyAmountAfterFee=TokenAmount(str(buy_amount)),
+    )
+    order_quote = asyncio.run(
+        order_book_api.post_quote(order_quote_request, order_side)
+    )
+    return Wei(order_quote.quote.sellAmount.root)
+
+
 @tenacity.retry(
     stop=tenacity.stop_after_attempt(3),
     wait=tenacity.wait_fixed(1),
     after=lambda x: logger.debug(f"get_buy_token_amount failed, {x.attempt_number=}."),
 )
 def get_buy_token_amount(
-    amount_wei: Wei,
+    sell_amount: Wei,
     sell_token: ChecksumAddress,
     buy_token: ChecksumAddress,
     chain: Chain = Chain.GNOSIS,
@@ -57,14 +95,20 @@ def get_buy_token_amount(
     )
     order_side = OrderQuoteSide1(
         kind=OrderQuoteSideKindSell.sell,
-        sellAmountBeforeFee=TokenAmount(str(amount_wei)),
+        sellAmountBeforeFee=TokenAmount(str(sell_amount)),
     )
     order_quote = asyncio.run(
         order_book_api.post_quote(order_quote_request, order_side)
     )
-    return wei_type(order_quote.quote.buyAmount.root)
+    return Wei(order_quote.quote.buyAmount.root)
 
 
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_fixed(1),
+    retry=tenacity.retry_if_not_exception_type((TimeoutError, OrderStatusError)),
+    after=lambda x: logger.debug(f"swap_tokens_waiting failed, {x.attempt_number=}."),
+)
 def swap_tokens_waiting(
     amount_wei: Wei,
     sell_token: ChecksumAddress,
@@ -102,7 +146,7 @@ async def swap_tokens_waiting_async(
     timeout: timedelta = timedelta(seconds=60),
 ) -> OrderMetaData:
     order = await swap_tokens(
-        amount=amount_wei,
+        amount=amount_wei.value,
         sell_token=sell_token,
         buy_token=buy_token,
         account=account,
@@ -125,7 +169,7 @@ async def swap_tokens_waiting_async(
             OrderStatus.cancelled,
             OrderStatus.expired,
         ):
-            raise ValueError(f"Order {order.uid} failed. {order.url}")
+            raise OrderStatusError(f"Order {order.uid} failed. {order.url}")
 
         if utcnow() - start_time > timeout:
             raise TimeoutError(

prediction_market_agent_tooling/tools/langfuse_client_utils.py

@@ -13,6 +13,9 @@ from prediction_market_agent_tooling.markets.data_models import (
     TradeType,
 )
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
+from prediction_market_agent_tooling.markets.omen.omen_constants import (
+    WRAPPED_XDAI_CONTRACT_ADDRESS,
+)
 from prediction_market_agent_tooling.tools.utils import DatetimeUTC
 
 
@@ -147,11 +150,20 @@ def get_trace_for_bet(
     # Filter for traces with the same bet outcome and amount
     traces_for_bet: list[ProcessMarketTrace] = []
     for t in traces:
+        if (
+            t.market.collateral_token_contract_address_checksummed
+            not in WRAPPED_XDAI_CONTRACT_ADDRESS
+        ):
+            # TODO: We need to compute bet amount token in USD here, but at the time of bet placement!
+            logger.warning(
+                "This currently works only for WXDAI markets, because we need to compare against USD value."
+            )
+            continue
         # Cannot use exact comparison due to gas fees
         if (
             t.buy_trade
             and t.buy_trade.outcome == bet.outcome
-            and np.isclose(t.buy_trade.amount.amount, bet.amount.amount)
+            and np.isclose(t.buy_trade.amount.value, bet.amount.value)
         ):
             traces_for_bet.append(t)
 

prediction_market_agent_tooling/tools/omen/sell_positions.py

@@ -1,6 +1,7 @@
 from datetime import timedelta
 
 from prediction_market_agent_tooling.config import APIKeys
+from prediction_market_agent_tooling.gtypes import USD
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
 from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import (
@@ -23,7 +24,9 @@ def sell_all(
         better_address=better_address,
         market_opening_after=utcnow() + timedelta(days=closing_later_than_days),
     )
-    bets_total_usd = sum(b.collateral_amount_usd for b in bets)
+    bets_total_usd = sum(
+        (b.get_collateral_amount_usd() for b in bets), start=USD.zero()
+    )
     unique_market_urls = set(b.fpmm.url for b in bets)
     starting_balance = get_balances(better_address)
     new_balance = starting_balance
@@ -37,9 +40,9 @@ def sell_all(
         outcome = agent_market.outcomes[bet.outcomeIndex]
         current_token_balance = agent_market.get_token_balance(better_address, outcome)
 
-        if current_token_balance.amount <= OmenAgentMarket.get_tiny_bet_amount().amount:
+        if current_token_balance.as_token <= agent_market.get_tiny_bet_amount():
             logger.info(
-                f"Skipping bet on {bet.fpmm.url} because the actual balance is unreasonably low {current_token_balance.amount}."
+                f"Skipping bet on {bet.fpmm.url} because the actual balance is unreasonably low {current_token_balance}."
             )
             continue
 

prediction_market_agent_tooling/tools/safe.py

@@ -11,11 +11,10 @@ from safe_eth.eth.constants import NULL_ADDRESS
 from safe_eth.eth.contracts import get_safe_V1_4_1_contract
 from safe_eth.safe.proxy_factory import ProxyFactoryV141
 from safe_eth.safe.safe import SafeV141
-from web3.types import Wei
 
+from prediction_market_agent_tooling.gtypes import Wei
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
-from prediction_market_agent_tooling.tools.web3_utils import wei_to_xdai
 
 
 def create_safe(
@@ -63,10 +62,10 @@ def create_safe(
             f"does not exist on network {ethereum_network.name}"
         )
 
-    account_balance = ethereum_client.get_balance(account.address)
-    account_balance_xdai = wei_to_xdai(account_balance)
+    account_balance = Wei(ethereum_client.get_balance(account.address))
+    account_balance_xdai = account_balance.as_token
     # We set a reasonable expected balance below for Safe deployment not to fail.
-    if account_balance_xdai < 0.01:
+    if account_balance_xdai.value < 0.01:
         raise ValueError(
             f"Client's balance is {account_balance_xdai} xDAI, too low for deploying a Safe."
         )
@@ -108,7 +107,7 @@ def create_safe(
             payment_token,
             payment,
             payment_receiver,
-        ).build_transaction({"gas": 1, "gasPrice": Wei(1)})["data"]
+        ).build_transaction({"gas": 1, "gasPrice": Wei(1).value})["data"]
     )
 
     proxy_factory = ProxyFactoryV141(proxy_factory_address, ethereum_client)