prediction-market-agent-tooling 0.61.1.dev463__py3-none-any.whl → 0.61.1.dev477__py3-none-any.whl

prediction_market_agent_tooling/markets/omen/omen.py

@@ -1,3 +1,4 @@
+import sys
 import typing as t
 from datetime import timedelta
 
@@ -6,17 +7,16 @@ from web3 import Web3
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.gtypes import (
-    USD,
     ChecksumAddress,
     HexAddress,
     HexStr,
+    OmenOutcomeToken,
     OutcomeStr,
-    OutcomeToken,
-    OutcomeWei,
     Probability,
-    Token,
     Wei,
+    wei_type,
     xDai,
+    xdai_type,
 )
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.agent_market import (
@@ -27,11 +27,17 @@ from prediction_market_agent_tooling.markets.agent_market import (
     ProcessedTradedMarket,
     SortBy,
 )
-from prediction_market_agent_tooling.markets.blockchain_utils import store_trades
+from prediction_market_agent_tooling.markets.blockchain_utils import (
+    get_total_balance,
+    store_trades,
+)
 from prediction_market_agent_tooling.markets.data_models import (
     Bet,
-    ExistingPosition,
+    BetAmount,
+    Currency,
+    Position,
     ResolvedBet,
+    TokenAmount,
 )
 from prediction_market_agent_tooling.markets.omen.data_models import (
     OMEN_FALSE_OUTCOME,
@@ -59,7 +65,6 @@ from prediction_market_agent_tooling.markets.omen.omen_contracts import (
     build_parent_collection_id,
 )
 from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import (
-    SAFE_COLLATERAL_TOKENS_ADDRESSES,
     OmenSubgraphHandler,
 )
 from prediction_market_agent_tooling.tools.balances import get_balances
@@ -76,20 +81,21 @@ from prediction_market_agent_tooling.tools.tokens.auto_withdraw import (
     auto_withdraw_collateral_token,
 )
 from prediction_market_agent_tooling.tools.tokens.main_token import KEEPING_ERC20_TOKEN
-from prediction_market_agent_tooling.tools.tokens.usd import (
-    get_token_in_usd,
-    get_usd_in_token,
-    get_xdai_in_usd,
-)
 from prediction_market_agent_tooling.tools.utils import (
     DatetimeUTC,
     calculate_sell_amount_in_collateral,
     check_not_none,
 )
-from prediction_market_agent_tooling.tools.web3_utils import get_receipt_block_timestamp
+from prediction_market_agent_tooling.tools.web3_utils import (
+    add_fraction,
+    get_receipt_block_timestamp,
+    remove_fraction,
+    wei_to_xdai,
+    xdai_to_wei,
+)
 
-OMEN_DEFAULT_REALITIO_BOND_VALUE = xDai(0.01)
-OMEN_TINY_BET_AMOUNT = USD(0.00001)
+OMEN_DEFAULT_REALITIO_BOND_VALUE = xdai_type(0.01)
+OMEN_TINY_BET_AMOUNT = xdai_type(0.00001)
 
 
 class OmenAgentMarket(AgentMarket):
@@ -97,6 +103,7 @@ class OmenAgentMarket(AgentMarket):
     Omen's market class that can be used by agents to make predictions.
     """
 
+    currency: t.ClassVar[Currency] = Currency.xDai
     base_url: t.ClassVar[str] = PRESAGIO_BASE_URL
     creator: HexAddress
 
@@ -145,24 +152,25 @@ class OmenAgentMarket(AgentMarket):
     def get_liquidity_in_wei(self, web3: Web3 | None = None) -> Wei:
         return self.get_contract().totalSupply(web3)
 
-    def get_liquidity(self, web3: Web3 | None = None) -> Token:
-        return self.get_liquidity_in_wei(web3).as_token
-
-    def get_tiny_bet_amount(self) -> Token:
-        return self.get_in_token(OMEN_TINY_BET_AMOUNT)
+    def get_liquidity_in_xdai(self, web3: Web3 | None = None) -> xDai:
+        return wei_to_xdai(self.get_liquidity_in_wei(web3))
 
-    def get_token_in_usd(self, x: Token) -> USD:
-        return get_token_in_usd(x, self.collateral_token_contract_address_checksummed)
+    def get_liquidity(self) -> TokenAmount:
+        return TokenAmount(
+            amount=self.get_liquidity_in_xdai(),
+            currency=self.currency,
+        )
 
-    def get_usd_in_token(self, x: USD) -> Token:
-        return get_usd_in_token(x, self.collateral_token_contract_address_checksummed)
+    @classmethod
+    def get_tiny_bet_amount(cls) -> BetAmount:
+        return BetAmount(amount=OMEN_TINY_BET_AMOUNT, currency=cls.currency)
 
     def liquidate_existing_positions(
         self,
         bet_outcome: bool,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
-        larger_than: OutcomeToken | None = None,
+        larger_than: float | None = None,
     ) -> None:
         """
         Liquidates all previously existing positions.
@@ -171,14 +179,16 @@ class OmenAgentMarket(AgentMarket):
         api_keys = api_keys if api_keys is not None else APIKeys()
         better_address = api_keys.bet_from_address
         larger_than = (
-            larger_than if larger_than is not None else self.get_liquidatable_amount()
+            larger_than
+            if larger_than is not None
+            else self.get_liquidatable_amount().amount
         )
         prev_positions_for_market = self.get_positions(
             user_id=better_address, liquid_only=True, larger_than=larger_than
         )
 
         for prev_position in prev_positions_for_market:
-            for position_outcome, token_amount in prev_position.amounts_ot.items():
+            for position_outcome, token_amount in prev_position.amounts.items():
                 position_outcome_bool = get_boolean_outcome(position_outcome)
                 if position_outcome_bool != bet_outcome:
                     self.sell_tokens(
@@ -192,7 +202,7 @@ class OmenAgentMarket(AgentMarket):
     def place_bet(
         self,
         outcome: bool,
-        amount: USD,
+        amount: BetAmount,
         auto_deposit: bool = True,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
@@ -201,9 +211,12 @@ class OmenAgentMarket(AgentMarket):
             raise ValueError(
                 f"Market {self.id} is not open for trading. Cannot place bet."
             )
+        if amount.currency != self.currency:
+            raise ValueError(f"Omen bets are made in xDai. Got {amount.currency}.")
+        amount_xdai = xDai(amount.amount)
         return binary_omen_buy_outcome_tx(
             api_keys=api_keys if api_keys is not None else APIKeys(),
-            amount=amount,
+            amount=amount_xdai,
             market=self,
             binary_outcome=outcome,
             auto_deposit=auto_deposit,
@@ -213,7 +226,7 @@ class OmenAgentMarket(AgentMarket):
     def buy_tokens(
         self,
         outcome: bool,
-        amount: USD,
+        amount: TokenAmount,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
     ) -> str:
@@ -224,37 +237,28 @@ class OmenAgentMarket(AgentMarket):
             api_keys=api_keys,
         )
 
-    def get_sell_value_of_outcome_token(
-        self, outcome: str, amount: OutcomeToken, web3: Web3 | None = None
-    ) -> Token:
-        """
-        Market can have as collateral token GNO for example.
-        When you place bet, you buy shares with GNO. For example, you get 10 shares for 1 GNO.
-        When selling, you need to provide the amount in GNO, which is cumbersome because you know how much shares you have, but you don't have the price of the shares in GNO.
-        Use this to convert how much collateral token (GNO in our example) to sell, to get the amount of shares you want to sell.
-        """
-        outcome_bool = get_boolean_outcome(outcome)
-
+    def calculate_sell_amount_in_collateral(
+        self, amount: TokenAmount, outcome: bool, web3: Web3 | None = None
+    ) -> xDai:
         pool_balance = get_conditional_tokens_balance_for_market(
             self, self.market_maker_contract_address_checksummed, web3=web3
         )
 
-        sell_str = self.outcomes[self.yes_index if outcome_bool else self.no_index]
-        other_str = self.outcomes[self.no_index if outcome_bool else self.yes_index]
+        sell_str = self.outcomes[self.yes_index if outcome else self.no_index]
+        other_str = self.outcomes[self.no_index if outcome else self.yes_index]
 
         collateral = calculate_sell_amount_in_collateral(
-            shares_to_sell=amount,
-            holdings=pool_balance[self.get_index_set(sell_str)].as_outcome_token,
-            other_holdings=pool_balance[self.get_index_set(other_str)].as_outcome_token,
+            shares_to_sell=amount.amount,
+            holdings=wei_to_xdai(pool_balance[self.get_index_set(sell_str)]),
+            other_holdings=wei_to_xdai(pool_balance[self.get_index_set(other_str)]),
             fees=self.fees,
         )
-
-        return collateral
+        return xDai(collateral)
 
     def sell_tokens(
         self,
         outcome: bool,
-        amount: USD | OutcomeToken,
+        amount: TokenAmount,
         auto_withdraw: bool = True,
         api_keys: APIKeys | None = None,
         web3: Web3 | None = None,
@@ -263,8 +267,16 @@ class OmenAgentMarket(AgentMarket):
             raise ValueError(
                 f"Market {self.id} is not open for trading. Cannot sell tokens."
             )
-        return binary_omen_sell_outcome_tx(
+
+        # Convert from token (i.e. share) number to xDai value of tokens, as
+        # this is the expected unit of the argument in the smart contract.
+        collateral = self.calculate_sell_amount_in_collateral(
             amount=amount,
+            outcome=outcome,
+            web3=web3,
+        )
+        return binary_omen_sell_outcome_tx(
+            amount=collateral,
             api_keys=api_keys if api_keys is not None else APIKeys(),
             market=self,
             binary_outcome=outcome,
@@ -312,7 +324,7 @@ class OmenAgentMarket(AgentMarket):
             better_address=user,
             position_id_in=[p.id for p in positions],
             # After redeem, this will became zero.
-            total_balance_bigger_than=OutcomeWei(0),
+            total_balance_bigger_than=wei_type(0),
         )
         return len(user_positions) > 0
 
@@ -349,16 +361,16 @@ class OmenAgentMarket(AgentMarket):
             current_p_yes=model.current_p_yes,
             condition=model.condition,
             url=model.url,
-            volume=model.collateralVolume.as_token,
+            volume=wei_to_xdai(model.collateralVolume),
             close_time=model.close_time,
             fees=MarketFees(
                 bet_proportion=(
-                    model.fee.as_token.value if model.fee is not None else 0.0
+                    float(wei_to_xdai(model.fee)) if model.fee is not None else 0.0
                 ),
                 absolute=0,
             ),
             outcome_token_pool={
-                model.outcomes[i]: model.outcomeTokenAmounts[i].as_outcome_token
+                model.outcomes[i]: wei_to_xdai(Wei(model.outcomeTokenAmounts[i]))
                 for i in range(len(model.outcomes))
             },
         )
@@ -395,24 +407,22 @@ class OmenAgentMarket(AgentMarket):
         redeem_from_all_user_positions(api_keys)
 
     @staticmethod
-    def get_trade_balance(api_keys: APIKeys, web3: Web3 | None = None) -> USD:
-        native_usd = get_xdai_in_usd(
-            get_balances(api_keys.bet_from_address, web3=web3).xdai
-        )
-        keeping_usd = get_token_in_usd(
-            KEEPING_ERC20_TOKEN.balance_of_in_tokens(
-                api_keys.bet_from_address, web3=web3
-            ),
-            KEEPING_ERC20_TOKEN.address,
+    def get_trade_balance(api_keys: APIKeys, web3: Web3 | None = None) -> xDai:
+        native_token_balance = get_balances(api_keys.bet_from_address, web3=web3).xdai
+        return xdai_type(
+            wei_to_xdai(
+                KEEPING_ERC20_TOKEN.balanceOf(api_keys.bet_from_address, web3=web3)
+            )
+            + native_token_balance
         )
-        return keeping_usd + native_usd
 
     @staticmethod
     def verify_operational_balance(api_keys: APIKeys) -> bool:
-        return get_balances(
-            # Use `public_key`, not `bet_from_address` because transaction costs are paid from the EOA wallet.
+        return get_total_balance(
             api_keys.public_key,
-        ).xdai > xDai(0.001)
+            # Use `public_key`, not `bet_from_address` because transaction costs are paid from the EOA wallet.
+            sum_wxdai=False,
+        ) > xdai_type(0.001)
 
     def store_prediction(
         self, processed_market: ProcessedMarket | None, keys: APIKeys, agent_name: str
@@ -487,19 +497,22 @@ class OmenAgentMarket(AgentMarket):
 
     def get_token_balance(
         self, user_id: str, outcome: str, web3: Web3 | None = None
-    ) -> OutcomeToken:
+    ) -> TokenAmount:
         index_set = self.get_index_set(outcome)
         balances = get_conditional_tokens_balance_for_market(
             self, Web3.to_checksum_address(user_id), web3=web3
         )
-        return balances[index_set].as_outcome_token
+        return TokenAmount(
+            amount=wei_to_xdai(balances[index_set]),
+            currency=self.currency,
+        )
 
-    def get_position(self, user_id: str) -> ExistingPosition | None:
+    def get_position(self, user_id: str) -> Position | None:
         liquidatable_amount = self.get_liquidatable_amount()
         existing_positions = self.get_positions(
             user_id=user_id,
             liquid_only=True,
-            larger_than=liquidatable_amount,
+            larger_than=liquidatable_amount.amount,
         )
         existing_position = next(
             iter([i for i in existing_positions if i.market_id == self.id]), None
@@ -511,12 +524,12 @@ class OmenAgentMarket(AgentMarket):
         cls,
         user_id: str,
         liquid_only: bool = False,
-        larger_than: OutcomeToken = OutcomeToken(0),
-    ) -> t.Sequence[ExistingPosition]:
+        larger_than: float = 0,
+    ) -> list[Position]:
         sgh = OmenSubgraphHandler()
         omen_positions = sgh.get_user_positions(
             better_address=Web3.to_checksum_address(user_id),
-            total_balance_bigger_than=larger_than.as_outcome_wei,
+            total_balance_bigger_than=xdai_to_wei(xDai(larger_than)),
         )
 
         # Sort positions and corresponding markets by condition_id
@@ -550,105 +563,142 @@ class OmenAgentMarket(AgentMarket):
             if liquid_only and not market.can_be_traded():
                 continue
 
-            amounts_ot: dict[OutcomeStr, OutcomeToken] = {}
-
+            amounts: dict[OutcomeStr, TokenAmount] = {}
             for omen_position in omen_positions:
                 outecome_str = market.index_set_to_outcome_str(
                     omen_position.position.index_set
                 )
 
                 # Validate that outcomes are unique for a given condition_id.
-                if outecome_str in amounts_ot:
+                if outecome_str in amounts:
                     raise ValueError(
-                        f"Outcome {outecome_str} already exists in {amounts_ot=}"
+                        f"Outcome {outecome_str} already exists in {amounts=}"
                     )
 
-                amounts_ot[outecome_str] = omen_position.totalBalance.as_outcome_token
-
-            amounts_current = {
-                k: market.get_token_in_usd(
-                    # If the market is not open for trading anymore, then current value is equal to potential value.
-                    market.get_sell_value_of_outcome_token(k, v)
-                    if market.can_be_traded()
-                    else v.as_token
-                )
-                for k, v in amounts_ot.items()
-            }
-            amounts_potential = {
-                k: market.get_token_in_usd(v.as_token) for k, v in amounts_ot.items()
-            }
-            positions.append(
-                ExistingPosition(
-                    market_id=market.id,
-                    amounts_current=amounts_current,
-                    amounts_potential=amounts_potential,
-                    amounts_ot=amounts_ot,
+                amounts[outecome_str] = TokenAmount(
+                    amount=wei_to_xdai(omen_position.totalBalance),
+                    currency=cls.currency,
                 )
-            )
+
+            positions.append(Position(market_id=market.id, amounts=amounts))
 
         return positions
 
+    @classmethod
+    def get_positions_value(cls, positions: list[Position]) -> BetAmount:
+        # Two dicts to map from market ids to (1) positions and (2) market.
+        market_ids_positions = {p.market_id: p for p in positions}
+        # Check there is only one position per market.
+        if len(set(market_ids_positions.keys())) != len(positions):
+            raise ValueError(
+                f"Markets for positions ({market_ids_positions.keys()}) are not unique."
+            )
+        markets: list[OmenAgentMarket] = [
+            OmenAgentMarket.from_data_model(m)
+            for m in OmenSubgraphHandler().get_omen_binary_markets(
+                limit=sys.maxsize, id_in=list(market_ids_positions.keys())
+            )
+        ]
+        market_ids_markets = {m.id: m for m in markets}
+
+        # Validate that dict keys are the same.
+        if set(market_ids_positions.keys()) != set(market_ids_markets.keys()):
+            raise ValueError(
+                f"Market ids in {market_ids_positions.keys()} are not the same as in {market_ids_markets.keys()}"
+            )
+
+        # Initialise position value.
+        total_position_value = 0.0
+
+        for market_id in market_ids_positions.keys():
+            position = market_ids_positions[market_id]
+            market = market_ids_markets[market_id]
+
+            yes_tokens = 0.0
+            no_tokens = 0.0
+            if OMEN_TRUE_OUTCOME in position.amounts:
+                yes_tokens = position.amounts[OutcomeStr(OMEN_TRUE_OUTCOME)].amount
+            if OMEN_FALSE_OUTCOME in position.amounts:
+                no_tokens = position.amounts[OutcomeStr(OMEN_FALSE_OUTCOME)].amount
+
+            # Account for the value of positions in resolved markets
+            if market.is_resolved() and market.has_successful_resolution():
+                valued_tokens = yes_tokens if market.boolean_outcome else no_tokens
+                total_position_value += valued_tokens
+
+            # Or if the market is open and trading, get the value of the position
+            elif market.can_be_traded():
+                total_position_value += yes_tokens * market.yes_outcome_price
+                total_position_value += no_tokens * market.no_outcome_price
+
+            # Or if the market is still open but not trading, estimate the value
+            # of the position
+            else:
+                if yes_tokens:
+                    yes_price = check_not_none(
+                        market.get_last_trade_yes_outcome_price()
+                    )
+                    total_position_value += yes_tokens * yes_price
+                if no_tokens:
+                    no_price = check_not_none(market.get_last_trade_no_outcome_price())
+                    total_position_value += no_tokens * no_price
+
+        return BetAmount(amount=total_position_value, currency=cls.currency)
+
     @classmethod
     def get_user_url(cls, keys: APIKeys) -> str:
         return get_omen_user_url(keys.bet_from_address)
 
     def get_buy_token_amount(
-        self, bet_amount: USD | Token, direction: bool
-    ) -> OutcomeToken:
+        self, bet_amount: BetAmount, direction: bool
+    ) -> TokenAmount:
         """
         Note: this is only valid if the market instance's token pool is
         up-to-date with the smart contract.
         """
         outcome_token_pool = check_not_none(self.outcome_token_pool)
         amount = get_buy_outcome_token_amount(
-            investment_amount=self.get_in_token(bet_amount),
+            investment_amount=bet_amount.amount,
             buy_direction=direction,
             yes_outcome_pool_size=outcome_token_pool[OMEN_TRUE_OUTCOME],
             no_outcome_pool_size=outcome_token_pool[OMEN_FALSE_OUTCOME],
             fees=self.fees,
         )
-        return amount
+        return TokenAmount(amount=amount, currency=self.currency)
 
     def _get_buy_token_amount_from_smart_contract(
-        self, bet_amount: USD, direction: bool
-    ) -> OutcomeToken:
-        bet_amount_in_tokens = get_usd_in_token(
-            bet_amount, self.collateral_token_contract_address_checksummed
-        )
-        received_token_amount_wei = self.get_contract().calcBuyAmount(
-            investment_amount=bet_amount_in_tokens.as_wei,
-            outcome_index=self.get_outcome_index(
-                self.get_outcome_str_from_bool(direction)
-            ),
+        self, bet_amount: BetAmount, direction: bool
+    ) -> TokenAmount:
+        received_token_amount_wei = Wei(
+            self.get_contract().calcBuyAmount(
+                investment_amount=xdai_to_wei(xDai(bet_amount.amount)),
+                outcome_index=self.get_outcome_index(
+                    self.get_outcome_str_from_bool(direction)
+                ),
+            )
         )
-        received_token_amount = received_token_amount_wei.as_outcome_token
-        return received_token_amount
+        received_token_amount = float(wei_to_xdai(received_token_amount_wei))
+        return TokenAmount(amount=received_token_amount, currency=self.currency)
 
-    def get_new_p_yes(self, bet_amount: USD, direction: bool) -> Probability:
+    def get_new_p_yes(self, bet_amount: BetAmount, direction: bool) -> Probability:
         """
         Calculate the new p_yes based on the bet amount and direction.
         """
         if not self.has_token_pool():
             raise ValueError("Outcome token pool is required to calculate new p_yes.")
 
-        bet_amount_in_tokens = self.get_usd_in_token(bet_amount)
         outcome_token_pool = check_not_none(self.outcome_token_pool)
-
-        yes_outcome_pool_size = outcome_token_pool[
-            self.get_outcome_str_from_bool(True)
-        ].value
-        no_outcome_pool_size = outcome_token_pool[
-            self.get_outcome_str_from_bool(False)
-        ].value
+        yes_outcome_pool_size = outcome_token_pool[self.get_outcome_str_from_bool(True)]
+        no_outcome_pool_size = outcome_token_pool[self.get_outcome_str_from_bool(False)]
 
         new_yes_outcome_pool_size = yes_outcome_pool_size + (
-            self.fees.get_after_fees(bet_amount_in_tokens).value
+            self.fees.get_bet_size_after_fees(bet_amount.amount)
         )
         new_no_outcome_pool_size = no_outcome_pool_size + (
-            self.fees.get_after_fees(bet_amount_in_tokens).value
+            self.fees.get_bet_size_after_fees(bet_amount.amount)
         )
 
-        received_token_amount = self.get_buy_token_amount(bet_amount, direction).value
+        received_token_amount = self.get_buy_token_amount(bet_amount, direction).amount
         if direction:
             new_yes_outcome_pool_size -= received_token_amount
         else:
@@ -687,18 +737,11 @@ def get_omen_user_url(address: ChecksumAddress) -> str:
 
 
 def pick_binary_market(
-    sort_by: SortBy = SortBy.CLOSING_SOONEST,
-    filter_by: FilterBy = FilterBy.OPEN,
-    collateral_token_address_in: (
-        tuple[ChecksumAddress, ...] | None
-    ) = SAFE_COLLATERAL_TOKENS_ADDRESSES,
+    sort_by: SortBy = SortBy.CLOSING_SOONEST, filter_by: FilterBy = FilterBy.OPEN
 ) -> OmenMarket:
     subgraph_handler = OmenSubgraphHandler()
     return subgraph_handler.get_omen_binary_markets_simple(
-        limit=1,
-        sort_by=sort_by,
-        filter_by=filter_by,
-        collateral_token_address_in=collateral_token_address_in,
+        limit=1, sort_by=sort_by, filter_by=filter_by
     )[0]
 
 
@@ -709,7 +752,7 @@ def pick_binary_market(
 )
 def omen_buy_outcome_tx(
     api_keys: APIKeys,
-    amount: USD | Token,
+    amount: xDai,
     market: OmenAgentMarket,
     outcome: str,
     auto_deposit: bool,
@@ -717,32 +760,35 @@ def omen_buy_outcome_tx(
     slippage: float = 0.01,
 ) -> str:
     """
-    Bets the given amount for the given outcome in the given market.
+    Bets the given amount of xDai for the given outcome in the given market.
     """
+    amount_wei = xdai_to_wei(amount)
+
     market_contract: OmenFixedProductMarketMakerContract = market.get_contract()
     collateral_token_contract = market_contract.get_collateral_token_contract(web3)
 
-    amount_token = market.get_in_token(amount)
-    amount_wei = amount_token.as_wei
+    # In case of ERC4626, obtained (for example) sDai out of xDai could be lower than the `amount_wei`, so we need to handle it.
+    amount_wei_to_buy = collateral_token_contract.get_in_shares(amount_wei, web3)
 
     # Get the index of the outcome we want to buy.
     outcome_index: int = market.get_outcome_index(outcome)
 
     # Calculate the amount of shares we will get for the given investment amount.
     expected_shares = market_contract.calcBuyAmount(
-        amount_wei, outcome_index, web3=web3
+        amount_wei_to_buy, outcome_index, web3=web3
     )
     # Allow small slippage.
-    expected_shares = expected_shares.without_fraction(slippage)
+    expected_shares = remove_fraction(expected_shares, slippage)
     # Approve the market maker to withdraw our collateral token.
     collateral_token_contract.approve(
         api_keys=api_keys,
         for_address=market_contract.address,
-        amount_wei=amount_wei,
+        amount_wei=amount_wei_to_buy,
         web3=web3,
     )
 
     if auto_deposit:
+        # In auto-depositing, we need to deposit the original `amount_wei`, e.g. we can deposit 2 xDai, but receive 1.8 sDai, so for the bet we will use `amount_wei_to_buy`.
         auto_deposit_collateral_token(
             collateral_token_contract, amount_wei, api_keys, web3
         )
@@ -750,7 +796,7 @@ def omen_buy_outcome_tx(
     # Buy shares using the deposited xDai in the collateral token.
     tx_receipt = market_contract.buy(
         api_keys=api_keys,
-        amount_wei=amount_wei,
+        amount_wei=amount_wei_to_buy,
         outcome_index=outcome_index,
         min_outcome_tokens_to_buy=expected_shares,
         web3=web3,
@@ -761,7 +807,7 @@ def omen_buy_outcome_tx(
 
 def binary_omen_buy_outcome_tx(
     api_keys: APIKeys,
-    amount: USD | Token,
+    amount: xDai,
     market: OmenAgentMarket,
     binary_outcome: bool,
     auto_deposit: bool,
@@ -779,7 +825,7 @@ def binary_omen_buy_outcome_tx(
 
 def omen_sell_outcome_tx(
     api_keys: APIKeys,
-    amount: OutcomeToken | Token | USD,
+    amount: xDai,  # The xDai value of shares to sell.
     market: OmenAgentMarket,
     outcome: str,
     auto_withdraw: bool,
@@ -793,18 +839,13 @@ def omen_sell_outcome_tx(
     The number of shares sold will depend on the share price at the time of the
     transaction.
     """
+    amount_wei = xdai_to_wei(amount)
+    amount_wei = remove_fraction(amount_wei, slippage)
+
     market_contract: OmenFixedProductMarketMakerContract = market.get_contract()
     conditional_token_contract = OmenConditionalTokenContract()
     collateral_token_contract = market_contract.get_collateral_token_contract(web3)
 
-    amount_token = (
-        market.get_sell_value_of_outcome_token(outcome, amount, web3)
-        if isinstance(amount, OutcomeToken)
-        else market.get_in_token(amount)
-    )
-    amount_wei = amount_token.as_wei
-    amount_wei = amount_wei.without_fraction(slippage)
-
     # Verify, that markets uses conditional tokens that we expect.
     if (
         market_contract.conditionalTokens(web3=web3)
@@ -814,15 +855,15 @@ def omen_sell_outcome_tx(
             f"Market {market.id} uses conditional token that we didn't expect, {market_contract.conditionalTokens()} != {conditional_token_contract.address=}"
         )
 
-    # Get the index of the outcome we want to sell.
+    # Get the index of the outcome we want to buy.
     outcome_index: int = market.get_outcome_index(outcome)
 
-    # Calculate the amount of shares we will sell for the given selling amount of collateral.
+    # Calculate the amount of shares we will sell for the given selling amount of xdai.
     max_outcome_tokens_to_sell = market_contract.calcSellAmount(
         amount_wei, outcome_index, web3=web3
     )
     # Allow small slippage.
-    max_outcome_tokens_to_sell = max_outcome_tokens_to_sell.with_fraction(slippage)
+    max_outcome_tokens_to_sell = add_fraction(max_outcome_tokens_to_sell, slippage)
 
     # Approve the market maker to move our (all) conditional tokens.
     conditional_token_contract.setApprovalForAll(
@@ -852,7 +893,7 @@ def omen_sell_outcome_tx(
 
 def binary_omen_sell_outcome_tx(
     api_keys: APIKeys,
-    amount: OutcomeToken | Token | USD,
+    amount: xDai,
     market: OmenAgentMarket,
     binary_outcome: bool,
     auto_withdraw: bool,
@@ -870,16 +911,16 @@ def binary_omen_sell_outcome_tx(
 
 def omen_create_market_tx(
     api_keys: APIKeys,
-    initial_funds: USD | Token,
+    initial_funds: xDai,
     question: str,
     closing_time: DatetimeUTC,
     category: str,
     language: str,
-    outcomes: t.Sequence[str],
+    outcomes: list[str],
     auto_deposit: bool,
     finalization_timeout: timedelta = REALITY_DEFAULT_FINALIZATION_TIMEOUT,
     fee_perc: float = OMEN_DEFAULT_MARKET_FEE_PERC,
-    distribution_hint: list[OutcomeWei] | None = None,
+    distribution_hint: list[OmenOutcomeToken] | None = None,
     collateral_token_address: ChecksumAddress = WrappedxDaiContract().address,
     arbitrator: Arbitrator = Arbitrator.KLEROS_31_JURORS_WITH_APPEAL,
     web3: Web3 | None = None,
@@ -890,12 +931,7 @@ def omen_create_market_tx(
     web3 = (
         web3 or OmenFixedProductMarketMakerFactoryContract.get_web3()
     )  # Default to Gnosis web3.
-    initial_funds_in_collateral = (
-        get_usd_in_token(initial_funds, collateral_token_address)
-        if isinstance(initial_funds, USD)
-        else initial_funds
-    )
-    initial_funds_in_collateral_wei = initial_funds_in_collateral.as_wei
+    initial_funds_wei = xdai_to_wei(initial_funds)
 
     realitio_contract = OmenRealitioContract()
     conditional_token_contract = OmenConditionalTokenContract()
@@ -919,9 +955,9 @@ def omen_create_market_tx(
 
     if auto_deposit:
         auto_deposit_collateral_token(
-            collateral_token_contract,
-            initial_funds_in_collateral_wei,
+            collateral_token_contract=collateral_token_contract,
             api_keys=api_keys,
+            amount_wei=initial_funds_wei,
             web3=web3,
         )
 
@@ -955,16 +991,21 @@ def omen_create_market_tx(
             web3=web3,
         )
 
+    # In case of ERC4626, obtained (for example) sDai out of xDai could be lower than the `amount_wei`, so we need to handle it.
+    initial_funds_in_shares = collateral_token_contract.get_in_shares(
+        amount=initial_funds_wei, web3=web3
+    )
+
     # Approve the market maker to withdraw our collateral token.
     collateral_token_contract.approve(
         api_keys=api_keys,
         for_address=factory_contract.address,
-        amount_wei=initial_funds_in_collateral_wei,
+        amount_wei=initial_funds_in_shares,
         web3=web3,
     )
 
     # Create the market.
-    fee = Token(fee_perc).as_wei
+    fee = xdai_to_wei(xdai_type(fee_perc))
     (
         market_event,
         funding_event,
@@ -974,7 +1015,7 @@ def omen_create_market_tx(
         condition_id=condition_id,
         fee=fee,
         distribution_hint=distribution_hint,
-        initial_funds_wei=initial_funds_in_collateral_wei,
+        initial_funds_wei=initial_funds_in_shares,
         collateral_token_address=collateral_token_contract.address,
         web3=web3,
     )
@@ -991,7 +1032,7 @@ def omen_create_market_tx(
         condition_id=condition_id,
         question_event=question_event,
         condition_event=cond_event,
-        initial_funds=initial_funds_in_collateral_wei,
+        initial_funds=initial_funds_wei,
         fee=fee,
         distribution_hint=distribution_hint,
     )
@@ -1000,28 +1041,27 @@ def omen_create_market_tx(
 def omen_fund_market_tx(
     api_keys: APIKeys,
     market: OmenAgentMarket,
-    funds: USD | Token,
+    funds: Wei,
     auto_deposit: bool,
     web3: Web3 | None = None,
 ) -> None:
-    funds_in_collateral = market.get_in_token(funds)
-    funds_in_collateral_wei = funds_in_collateral.as_wei
     market_contract = market.get_contract()
     collateral_token_contract = market_contract.get_collateral_token_contract(web3=web3)
 
+    amount_to_fund = collateral_token_contract.get_in_shares(funds, web3)
+
     collateral_token_contract.approve(
         api_keys=api_keys,
         for_address=market_contract.address,
-        amount_wei=funds_in_collateral_wei,
+        amount_wei=amount_to_fund,
         web3=web3,
     )
 
     if auto_deposit:
-        auto_deposit_collateral_token(
-            collateral_token_contract, funds_in_collateral_wei, api_keys, web3
-        )
+        # In auto-depositing, we need to deposit the original `funds`, e.g. we can deposit 2 xDai, but receive 1.8 sDai, so for the funding we will use `amount_to_fund`.
+        auto_deposit_collateral_token(collateral_token_contract, funds, api_keys, web3)
 
-    market_contract.addFunding(api_keys, funds_in_collateral_wei, web3=web3)
+    market_contract.addFunding(api_keys, amount_to_fund, web3=web3)
 
 
 def omen_redeem_full_position_tx(
@@ -1054,7 +1094,7 @@ def omen_redeem_full_position_tx(
     amount_per_index = get_conditional_tokens_balance_for_market(
         market, from_address, web3
     )
-    amount_wei = sum(amount_per_index.values(), start=OutcomeWei.zero())
+    amount_wei = sum(amount_per_index.values())
     if amount_wei == 0:
         logger.debug("No balance to claim. Exiting.")
         return
@@ -1072,10 +1112,10 @@ def omen_redeem_full_position_tx(
         web3=web3,
     )
     new_balance = collateral_token_contract.balanceOf(from_address, web3=web3)
-    balance_diff = new_balance - original_balance
+    balance_diff = wei_type(new_balance - original_balance)
 
     logger.info(
-        f"Redeemed {balance_diff.as_token} {collateral_token_contract.symbol_cached(web3=web3)} from market {market.question=} ({market.url})."
+        f"Redeemed {wei_to_xdai(balance_diff)} {collateral_token_contract.symbol_cached(web3=web3)} from market {market.question=} ({market.url})."
     )
 
     if auto_withdraw:
@@ -1091,12 +1131,12 @@ def get_conditional_tokens_balance_for_market(
     market: OmenAgentMarket,
     from_address: ChecksumAddress,
     web3: Web3 | None = None,
-) -> dict[int, OutcomeWei]:
+) -> dict[int, Wei]:
     """
     We derive the withdrawable balance from the ConditionalTokens contract through CollectionId -> PositionId (which
     also serves as tokenId) -> TokenBalances.
     """
-    balance_per_index_set: dict[int, OutcomeWei] = {}
+    balance_per_index_set: dict[int, Wei] = {}
     conditional_token_contract = OmenConditionalTokenContract()
     parent_collection_id = build_parent_collection_id()
 
@@ -1113,7 +1153,7 @@ def get_conditional_tokens_balance_for_market(
         balance_for_position = conditional_token_contract.balanceOf(
             from_address=from_address, position_id=position_id, web3=web3
         )
-        balance_per_index_set[index_set] = balance_for_position
+        balance_per_index_set[index_set] = wei_type(balance_for_position)
 
     return balance_per_index_set
 
@@ -1174,11 +1214,11 @@ def omen_remove_fund_market_tx(
     )
 
     new_balance = market_collateral_token_contract.balanceOf(from_address, web3=web3)
-    balance_diff = new_balance - original_balance
+    balance_diff = wei_type(new_balance - original_balance)
 
     logger.debug(f"Result from merge positions {result}")
     logger.info(
-        f"Withdrawn {balance_diff.as_token} {market_collateral_token_contract.symbol_cached(web3=web3)} from liquidity at {market.url=}."
+        f"Withdrawn {wei_to_xdai(balance_diff)} {market_collateral_token_contract.symbol_cached(web3=web3)} from liquidity at {market.url=}."
     )
 
     if auto_withdraw:
@@ -1204,7 +1244,7 @@ def redeem_from_all_user_positions(
     user_positions = OmenSubgraphHandler().get_user_positions(
         public_key,
         # After redeem, this will became zero and we won't re-process it.
-        total_balance_bigger_than=OutcomeWei(0),
+        total_balance_bigger_than=wei_type(0),
     )
 
     for index, user_position in enumerate(user_positions):
@@ -1232,10 +1272,10 @@ def redeem_from_all_user_positions(
             web3=web3,
         )
         new_balance = collateral_token_contract.balanceOf(public_key, web3=web3)
-        balance_diff = new_balance - original_balance
+        balance_diff = wei_type(new_balance - original_balance)
 
         logger.info(
-            f"Redeemed {balance_diff.as_token} {collateral_token_contract.symbol_cached(web3=web3)} from position {user_position.id=}."
+            f"Redeemed {wei_to_xdai(balance_diff)} {collateral_token_contract.symbol_cached(web3=web3)} from position {user_position.id=}."
         )
 
         if auto_withdraw:
@@ -1306,21 +1346,23 @@ def send_keeping_token_to_eoa_xdai(
         )
         return
 
-    need_to_withdraw = (min_required_balance - current_balances_eoa.xdai) * multiplier
-    need_to_withdraw_wei = need_to_withdraw.as_xdai_wei
+    need_to_withdraw = xDai(
+        (min_required_balance - current_balances_eoa.xdai) * multiplier
+    )
+    need_to_withdraw_wei = xdai_to_wei(need_to_withdraw)
 
-    if current_balances_eoa.wxdai >= need_to_withdraw.as_token:
+    if current_balances_eoa.wxdai >= need_to_withdraw:
         # If EOA has enough of wxDai, simply withdraw it.
         logger.info(
             f"Withdrawing {need_to_withdraw} wxDai from EOA to keep the EOA's xDai balance above the minimum required balance {min_required_balance}."
         )
         wxdai_contract.withdraw(
             api_keys=api_keys.copy_without_safe_address(),
-            amount_wei=need_to_withdraw_wei.as_wei,
+            amount_wei=need_to_withdraw_wei,
             web3=web3,
         )
 
-    elif current_balances_betting.wxdai >= need_to_withdraw.as_token:
+    elif current_balances_betting.wxdai >= need_to_withdraw:
         # If Safe has enough of wxDai:
         # First send them to EOA's address.
         logger.info(
@@ -1330,7 +1372,7 @@ def send_keeping_token_to_eoa_xdai(
             api_keys=api_keys,
             sender=api_keys.bet_from_address,
             recipient=api_keys.public_key,
-            amount_wei=need_to_withdraw_wei.as_wei,
+            amount_wei=need_to_withdraw_wei,
             web3=web3,
         )
         # And then simply withdraw it.
@@ -1339,7 +1381,7 @@ def send_keeping_token_to_eoa_xdai(
         )
         wxdai_contract.withdraw(
             api_keys=api_keys.copy_without_safe_address(),
-            amount_wei=need_to_withdraw_wei.as_wei,
+            amount_wei=need_to_withdraw_wei,
             web3=web3,
         )
 
@@ -1350,37 +1392,30 @@ def send_keeping_token_to_eoa_xdai(
 
 
 def get_buy_outcome_token_amount(
-    investment_amount: Token,
+    investment_amount: float,
     buy_direction: bool,
-    yes_outcome_pool_size: OutcomeToken,
-    no_outcome_pool_size: OutcomeToken,
+    yes_outcome_pool_size: float,
+    no_outcome_pool_size: float,
     fees: MarketFees,
-) -> OutcomeToken:
+) -> float:
     """
     Calculates the amount of outcome tokens received for a given investment
 
     Taken from https://github.com/gnosis/conditional-tokens-market-makers/blob/6814c0247c745680bb13298d4f0dd7f5b574d0db/contracts/FixedProductMarketMaker.sol#L264
     """
-    investment_amount_minus_fees = fees.get_after_fees(investment_amount)
-    investment_amount_minus_fees_as_ot = OutcomeToken(
-        investment_amount_minus_fees.value
-    )
+    investment_amount_minus_fees = fees.get_bet_size_after_fees(investment_amount)
     buy_token_pool_balance = (
         yes_outcome_pool_size if buy_direction else no_outcome_pool_size
     )
 
     pool_balance = no_outcome_pool_size if buy_direction else yes_outcome_pool_size
-    denominator = pool_balance + investment_amount_minus_fees_as_ot
-    ending_outcome_balance = OutcomeToken(
-        buy_token_pool_balance * pool_balance / denominator
-    )
+    denominator = pool_balance + investment_amount_minus_fees
+    ending_outcome_balance = buy_token_pool_balance * pool_balance / denominator
 
     if ending_outcome_balance <= 0:
         raise ValueError("must have non-zero balances")
 
     result = (
-        buy_token_pool_balance
-        + investment_amount_minus_fees_as_ot
-        - ending_outcome_balance
+        buy_token_pool_balance + investment_amount_minus_fees - ending_outcome_balance
     )
     return result