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

prediction_market_agent_tooling/tools/balances.py

@@ -1,32 +1,30 @@
 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, Token, 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: Token
+    sdai: Token
 
     @property
-    def total(self) -> xDai:
-        return xDai(self.xdai + self.wxdai + self.sdai)
+    def total(self) -> Token:
+        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 OutcomeToken, Token
 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: Token,
     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 = Token(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: Token,
     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=Token(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=Token(min(max_bet.value, abs(kelly_bet_amount))),
     )

prediction_market_agent_tooling/tools/betting_strategies/market_moving.py

@@ -2,19 +2,18 @@ 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 OutcomeToken, Probability, Token
 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 +41,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 = Token(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 +63,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 +102,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 Token
+
 
 class SimpleBet(BaseModel):
     direction: bool
-    size: float
+    size: Token

prediction_market_agent_tooling/tools/contract.py

@@ -15,6 +15,7 @@ from prediction_market_agent_tooling.gtypes import (
     ChainID,
     ChecksumAddress,
     Nonce,
+    Token,
     TxParams,
     TxReceipt,
     Wei,
@@ -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
+    ) -> Token:
+        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, Token, 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: Token,
         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

@@ -21,10 +21,9 @@ 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
@@ -62,7 +61,7 @@ def get_buy_token_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)
 
 
 def swap_tokens_waiting(
@@ -102,7 +101,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,

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,
@@ -9,32 +9,44 @@ from prediction_market_agent_tooling.tools.contract import (
     ContractERC20OnGnosisChain,
     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.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,
+    slippage: 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 later changes in the price.
+    collateral_amount_wei = collateral_amount_wei.with_fraction(slippage)
+
     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 +67,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 +84,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 +101,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,32 +120,22 @@ 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
-    )
+    amount_to_sell_wei = remaining_to_get_in_collateral_wei * 1.01
     # If we don't have enough of the source token.
     if amount_to_sell_wei > ContractERC20OnGnosisChain(
         address=KEEPING_ERC20_TOKEN.address

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(
@@ -27,10 +20,7 @@ def auto_withdraw_collateral_token(
     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,
-    )
+    amount_wei = amount_wei.without_fraction(slippage)
 
     if not amount_wei:
         logger.warning(
@@ -51,7 +41,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 +50,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,