prediction-market-agent-tooling 0.61.1.dev489__py3-none-any.whl → 0.62.0__py3-none-any.whl

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)

prediction_market_agent_tooling/tools/tokens/auto_deposit.py

@@ -1,7 +1,7 @@
 from web3 import Web3
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.gtypes import Wei, wei_type
+from prediction_market_agent_tooling.gtypes import USD, Wei
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.contract import (
     ContractDepositableWrapperERC20BaseClass,
@@ -10,31 +10,46 @@ from prediction_market_agent_tooling.tools.contract import (
     ContractERC4626BaseClass,
 )
 from prediction_market_agent_tooling.tools.cow.cow_order import (
-    get_buy_token_amount,
+    get_sell_token_amount,
     swap_tokens_waiting,
 )
 from prediction_market_agent_tooling.tools.tokens.main_token import KEEPING_ERC20_TOKEN
+from prediction_market_agent_tooling.tools.tokens.usd import get_usd_in_token
 from prediction_market_agent_tooling.tools.utils import should_not_happen
-from prediction_market_agent_tooling.tools.web3_utils import wei_to_xdai
 
 
 def auto_deposit_collateral_token(
     collateral_token_contract: ContractERC20BaseClass,
-    amount_wei: Wei,
+    collateral_amount_wei_or_usd: Wei | USD,
     api_keys: APIKeys,
     web3: Web3 | None = None,
+    surplus: float = 0.01,
 ) -> None:
+    collateral_amount_wei = (
+        collateral_amount_wei_or_usd
+        if isinstance(collateral_amount_wei_or_usd, Wei)
+        else get_usd_in_token(
+            collateral_amount_wei_or_usd, collateral_token_contract.address
+        ).as_wei
+    )
+    # Deposit a bit more to cover any small changes in conversions.
+    collateral_amount_wei = collateral_amount_wei.with_fraction(surplus)
+
     if isinstance(collateral_token_contract, ContractDepositableWrapperERC20BaseClass):
         # In this case, we can use deposit function directly, no need to go through DEX.
         auto_deposit_depositable_wrapper_erc20(
-            collateral_token_contract, amount_wei, api_keys, web3
+            collateral_token_contract, collateral_amount_wei, api_keys, web3
         )
 
     elif isinstance(collateral_token_contract, ContractERC4626BaseClass):
-        auto_deposit_erc4626(collateral_token_contract, amount_wei, api_keys, web3)
+        auto_deposit_erc4626(
+            collateral_token_contract, collateral_amount_wei, api_keys, web3
+        )
 
     elif isinstance(collateral_token_contract, ContractERC20BaseClass):
-        auto_deposit_erc20(collateral_token_contract, amount_wei, api_keys, web3)
+        auto_deposit_erc20(
+            collateral_token_contract, collateral_amount_wei, api_keys, web3
+        )
 
     else:
         should_not_happen("Unsupported ERC20 contract type.")
@@ -55,16 +70,16 @@ def auto_deposit_depositable_wrapper_erc20(
         return
 
     # If we don't have enough, we need to deposit the difference.
-    left_to_deposit = Wei(amount_wei - collateral_token_balance)
+    left_to_deposit = amount_wei - collateral_token_balance
     logger.info(
-        f"Depositing {wei_to_xdai(left_to_deposit)} {collateral_token_contract.symbol()}."
+        f"Depositing {left_to_deposit.as_token} {collateral_token_contract.symbol()}."
     )
     collateral_token_contract.deposit(api_keys, left_to_deposit, web3=web3)
 
 
 def auto_deposit_erc4626(
     collateral_token_contract: ContractERC4626BaseClass,
-    asset_amount_wei: Wei,
+    collateral_amount_wei: Wei,
     api_keys: APIKeys,
     web3: Web3 | None,
 ) -> None:
@@ -72,12 +87,12 @@ def auto_deposit_erc4626(
     collateral_token_balance_in_shares = collateral_token_contract.balanceOf(
         for_address=for_address, web3=web3
     )
-    asset_amount_wei_in_shares = collateral_token_contract.convertToShares(
-        asset_amount_wei, web3
+    asset_amount_wei = collateral_token_contract.convertToAssets(
+        collateral_amount_wei, web3
     )
 
     # If we have enough shares, we don't need to deposit.
-    if collateral_token_balance_in_shares >= asset_amount_wei_in_shares:
+    if collateral_token_balance_in_shares >= collateral_amount_wei:
         return
 
     # If we need to deposit into erc4626, we first need to have enough of the asset token.
@@ -89,7 +104,7 @@ def auto_deposit_erc4626(
         collateral_token_balance_in_assets = collateral_token_contract.convertToAssets(
             collateral_token_balance_in_shares, web3
         )
-        left_to_deposit = Wei(asset_amount_wei - collateral_token_balance_in_assets)
+        left_to_deposit = asset_amount_wei - collateral_token_balance_in_assets
         if (
             collateral_token_contract.get_asset_token_balance(for_address, web3)
             < left_to_deposit
@@ -108,31 +123,25 @@ def auto_deposit_erc4626(
 
 def auto_deposit_erc20(
     collateral_token_contract: ContractERC20BaseClass,
-    amount_xdai_wei: Wei,
+    collateral_amount_wei: Wei,
     api_keys: APIKeys,
     web3: Web3 | None,
 ) -> None:
-    # How much it is in the other token (collateral token).
-    collateral_amount_wei = get_buy_token_amount(
-        amount_xdai_wei,
-        KEEPING_ERC20_TOKEN.address,
-        collateral_token_contract.address,
-    )
     # How much do we have already in the other token (collateral token).
     collateral_balance_wei = collateral_token_contract.balanceOf(
         api_keys.bet_from_address
     )
     # Amount of collateral token remaining to get.
     remaining_to_get_in_collateral_wei = max(
-        0, collateral_amount_wei - collateral_balance_wei
+        Wei(0), collateral_amount_wei - collateral_balance_wei
     )
     if not remaining_to_get_in_collateral_wei:
         return
-    # Estimate of how much of the source token we need to sell in order to fill the remaining collateral amount, with 1% slippage to be sure.
-    amount_to_sell_wei = wei_type(
-        (remaining_to_get_in_collateral_wei * amount_xdai_wei)
-        / collateral_amount_wei
-        * 1.01
+    # Get of how much of the source token we need to sell in order to fill the remaining collateral amount.
+    amount_to_sell_wei = get_sell_token_amount(
+        remaining_to_get_in_collateral_wei,
+        sell_token=KEEPING_ERC20_TOKEN.address,
+        buy_token=collateral_token_contract.address,
     )
     # If we don't have enough of the source token.
     if amount_to_sell_wei > ContractERC20OnGnosisChain(

prediction_market_agent_tooling/tools/tokens/auto_withdraw.py

@@ -7,16 +7,9 @@ from prediction_market_agent_tooling.tools.contract import (
     ContractERC20BaseClass,
     ContractERC4626BaseClass,
 )
-from prediction_market_agent_tooling.tools.cow.cow_order import (
-    get_buy_token_amount,
-    swap_tokens_waiting,
-)
+from prediction_market_agent_tooling.tools.cow.cow_order import swap_tokens_waiting
 from prediction_market_agent_tooling.tools.tokens.main_token import KEEPING_ERC20_TOKEN
 from prediction_market_agent_tooling.tools.utils import should_not_happen
-from prediction_market_agent_tooling.tools.web3_utils import (
-    remove_fraction,
-    wei_to_xdai,
-)
 
 
 def auto_withdraw_collateral_token(
@@ -24,14 +17,7 @@ def auto_withdraw_collateral_token(
     amount_wei: Wei,
     api_keys: APIKeys,
     web3: Web3 | None = None,
-    slippage: float = 0.001,
 ) -> None:
-    # Small slippage as exact exchange rate constantly changes and we don't care about small differences.
-    amount_wei = remove_fraction(
-        amount_wei,
-        slippage,
-    )
-
     if not amount_wei:
         logger.warning(
             f"Amount to withdraw is zero, skipping withdrawal of {collateral_token_contract.symbol_cached(web3)}."
@@ -51,7 +37,7 @@ def auto_withdraw_collateral_token(
     ):
         # If the ERC4626 is backed by KEEPING_ERC20_TOKEN, we can withdraw it directly, no need to go through DEX.
         logger.info(
-            f"Withdrawing {wei_to_xdai(amount_wei)} from {collateral_token_contract.symbol_cached(web3)} into {KEEPING_ERC20_TOKEN.symbol_cached(web3)}"
+            f"Withdrawing {amount_wei.as_token} from {collateral_token_contract.symbol_cached(web3)} into {KEEPING_ERC20_TOKEN.symbol_cached(web3)}"
         )
         collateral_token_contract.withdraw_in_shares(
             api_keys,
@@ -60,18 +46,11 @@ def auto_withdraw_collateral_token(
         )
     elif isinstance(collateral_token_contract, ContractERC20BaseClass):
         logger.info(
-            f"Swapping {wei_to_xdai(amount_wei)} {collateral_token_contract.symbol_cached(web3)} into {KEEPING_ERC20_TOKEN.symbol_cached(web3)}"
+            f"Swapping {amount_wei.as_token} from {collateral_token_contract.symbol_cached(web3)} into {KEEPING_ERC20_TOKEN.symbol_cached(web3)}"
         )
         # Otherwise, DEX will handle the rest of token swaps.
-        # First, convert `amount_wei` from xDai-based value into the collateral token-based value.
-        collateral_amount_wei = get_buy_token_amount(
-            amount_wei,
-            KEEPING_ERC20_TOKEN.address,
-            collateral_token_contract.address,
-        )
-        # And then sell it.
         swap_tokens_waiting(
-            amount_wei=collateral_amount_wei,
+            amount_wei=amount_wei,
             sell_token=collateral_token_contract.address,
             buy_token=KEEPING_ERC20_TOKEN.address,
             api_keys=api_keys,

prediction_market_agent_tooling/tools/tokens/main_token.py

@@ -1,4 +1,4 @@
-from prediction_market_agent_tooling.gtypes import xdai_type
+from prediction_market_agent_tooling.gtypes import xDai
 from prediction_market_agent_tooling.markets.omen.omen_constants import (
     WRAPPED_XDAI_CONTRACT_ADDRESS,
 )
@@ -15,4 +15,4 @@ KEEPING_ERC20_TOKEN = ContractDepositableWrapperERC20OnGnosisChain(
     address=WRAPPED_XDAI_CONTRACT_ADDRESS
 )
 
-MINIMUM_NATIVE_TOKEN_IN_EOA_FOR_FEES = xdai_type(0.1)
+MINIMUM_NATIVE_TOKEN_IN_EOA_FOR_FEES = xDai(0.1)

prediction_market_agent_tooling/tools/tokens/token_utils.py

@@ -0,0 +1,46 @@
+from eth_typing.evm import ChecksumAddress
+from web3 import Web3
+
+from prediction_market_agent_tooling.gtypes import ChecksumAddress, Wei
+from prediction_market_agent_tooling.tools.contract import (
+    ContractERC4626BaseClass,
+    init_collateral_token_contract,
+    to_gnosis_chain_contract,
+)
+from prediction_market_agent_tooling.tools.cow.cow_order import get_buy_token_amount
+
+
+def convert_to_another_token(
+    amount: Wei,
+    from_token: ChecksumAddress,
+    to_token: ChecksumAddress,
+    web3: Web3 | None = None,
+) -> Wei:
+    from_token_contract = to_gnosis_chain_contract(
+        init_collateral_token_contract(from_token, web3)
+    )
+    to_token_contract = to_gnosis_chain_contract(
+        init_collateral_token_contract(to_token, web3)
+    )
+
+    if from_token == to_token:
+        return amount
+
+    elif (
+        isinstance(to_token_contract, ContractERC4626BaseClass)
+        and to_token_contract.get_asset_token_contract().address == from_token
+    ):
+        return to_token_contract.convertToShares(amount)
+
+    elif (
+        isinstance(from_token_contract, ContractERC4626BaseClass)
+        and from_token_contract.get_asset_token_contract().address == to_token
+    ):
+        return from_token_contract.convertToAssets(amount)
+
+    else:
+        return get_buy_token_amount(
+            amount,
+            from_token,
+            to_token,
+        )

prediction_market_agent_tooling/tools/tokens/usd.py

@@ -0,0 +1,79 @@
+from cachetools import TTLCache, cached
+from eth_typing.evm import ChecksumAddress
+
+from prediction_market_agent_tooling.gtypes import (
+    USD,
+    ChecksumAddress,
+    CollateralToken,
+    xDai,
+)
+from prediction_market_agent_tooling.markets.omen.omen_constants import (
+    SDAI_CONTRACT_ADDRESS,
+    WRAPPED_XDAI_CONTRACT_ADDRESS,
+)
+from prediction_market_agent_tooling.tools.contract import ContractERC4626OnGnosisChain
+from prediction_market_agent_tooling.tools.cow.cow_order import get_buy_token_amount
+
+
+def get_usd_in_xdai(amount: USD) -> xDai:
+    # xDai is stable coin against USD, so for simplicity we just cast it.
+    return xDai(amount.value)
+
+
+def get_xdai_in_usd(amount: xDai) -> USD:
+    # xDai is stable coin against USD, so for simplicity we just cast it.
+    return USD(amount.value)
+
+
+def get_usd_in_token(amount: USD, token_address: ChecksumAddress) -> CollateralToken:
+    rate = get_single_usd_to_token_rate(token_address)
+    return CollateralToken(amount.value * rate.value)
+
+
+def get_token_in_usd(amount: CollateralToken, token_address: ChecksumAddress) -> USD:
+    rate = get_single_token_to_usd_rate(token_address)
+    return USD(amount.value * rate.value)
+
+
+# A short cache to not spam CoW and prevent timeouts, but still have relatively fresh data.
+@cached(TTLCache(maxsize=100, ttl=5 * 60))
+def get_single_token_to_usd_rate(token_address: ChecksumAddress) -> USD:
+    # (w)xDai is a stable coin against USD, so use it to estimate USD worth.
+    if WRAPPED_XDAI_CONTRACT_ADDRESS == token_address:
+        return USD(1.0)
+    # sDai is ERC4626 with wxDai as asset, we can take the rate directly from there instead of calling CoW.
+    if SDAI_CONTRACT_ADDRESS == token_address:
+        return USD(
+            ContractERC4626OnGnosisChain(address=SDAI_CONTRACT_ADDRESS)
+            .convertToAssets(CollateralToken(1).as_wei)
+            .as_token.value
+        )
+    in_wei = get_buy_token_amount(
+        sell_amount=CollateralToken(1).as_wei,
+        sell_token=token_address,
+        buy_token=WRAPPED_XDAI_CONTRACT_ADDRESS,
+    )
+    in_token = in_wei.as_token
+    return USD(in_token.value)
+
+
+# A short cache to not spam CoW and prevent timeouts, but still have relatively fresh data.
+@cached(TTLCache(maxsize=100, ttl=5 * 60))
+def get_single_usd_to_token_rate(token_address: ChecksumAddress) -> CollateralToken:
+    # (w)xDai is a stable coin against USD, so use it to estimate USD worth.
+    if WRAPPED_XDAI_CONTRACT_ADDRESS == token_address:
+        return CollateralToken(1.0)
+    # sDai is ERC4626 with wxDai as asset, we can take the rate directly from there instead of calling CoW.
+    if SDAI_CONTRACT_ADDRESS == token_address:
+        return CollateralToken(
+            ContractERC4626OnGnosisChain(address=SDAI_CONTRACT_ADDRESS)
+            .convertToShares(CollateralToken(1).as_wei)
+            .as_token.value
+        )
+    in_wei = get_buy_token_amount(
+        sell_amount=CollateralToken(1).as_wei,
+        sell_token=WRAPPED_XDAI_CONTRACT_ADDRESS,
+        buy_token=token_address,
+    )
+    in_token = in_wei.as_token
+    return CollateralToken(in_token.value)

prediction_market_agent_tooling/tools/utils.py

@@ -9,7 +9,13 @@ from pydantic import BaseModel, ValidationError
 from scipy.optimize import newton
 from scipy.stats import entropy
 
-from prediction_market_agent_tooling.gtypes import DatetimeUTC, Probability, SecretStr
+from prediction_market_agent_tooling.gtypes import (
+    CollateralToken,
+    DatetimeUTC,
+    OutcomeToken,
+    Probability,
+    SecretStr,
+)
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.market_fees import MarketFees
 
@@ -182,11 +188,11 @@ def prob_uncertainty(prob: Probability) -> float:
 
 
 def calculate_sell_amount_in_collateral(
-    shares_to_sell: float,
-    holdings: float,
-    other_holdings: float,
+    shares_to_sell: OutcomeToken,
+    holdings: OutcomeToken,
+    other_holdings: OutcomeToken,
     fees: MarketFees,
-) -> float:
+) -> CollateralToken:
     """
     Computes the amount of collateral that needs to be sold to get `shares`
     amount of shares. Returns None if the amount can't be computed.
@@ -199,11 +205,11 @@ def calculate_sell_amount_in_collateral(
             raise ValueError("All share args must be greater than 0")
 
     def f(r: float) -> float:
-        R = (r + fees.absolute) / (1 - fees.bet_proportion)
+        R = OutcomeToken((r + fees.absolute) / (1 - fees.bet_proportion))
         first_term = other_holdings - R
         second_term = holdings + shares_to_sell - R
         third_term = holdings * other_holdings
-        return (first_term * second_term) - third_term
+        return ((first_term * second_term) - third_term).value
 
     amount_to_sell = newton(f, 0)
-    return float(amount_to_sell) * 0.999999  # Avoid rounding errors
+    return CollateralToken(float(amount_to_sell) * 0.999999)  # Avoid rounding errors