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

prediction_market_agent_tooling/markets/omen/omen.py

@@ -1,4 +1,3 @@
-import sys
 import typing as t
 from datetime import timedelta
 
@@ -7,16 +6,17 @@ 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,17 +27,11 @@ from prediction_market_agent_tooling.markets.agent_market import (
     ProcessedTradedMarket,
     SortBy,
 )
-from prediction_market_agent_tooling.markets.blockchain_utils import (
-    get_total_balance,
-    store_trades,
-)
+from prediction_market_agent_tooling.markets.blockchain_utils import store_trades
 from prediction_market_agent_tooling.markets.data_models import (
     Bet,
-    BetAmount,
-    Currency,
-    Position,
+    ExistingPosition,
     ResolvedBet,
-    TokenAmount,
 )
 from prediction_market_agent_tooling.markets.omen.data_models import (
     OMEN_FALSE_OUTCOME,
@@ -65,6 +59,7 @@ 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
@@ -81,21 +76,20 @@ 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 (
-    add_fraction,
-    get_receipt_block_timestamp,
-    remove_fraction,
-    wei_to_xdai,
-    xdai_to_wei,
-)
+from prediction_market_agent_tooling.tools.web3_utils import get_receipt_block_timestamp
 
-OMEN_DEFAULT_REALITIO_BOND_VALUE = xdai_type(0.01)
-OMEN_TINY_BET_AMOUNT = xdai_type(0.00001)
+OMEN_DEFAULT_REALITIO_BOND_VALUE = xDai(0.01)
+OMEN_TINY_BET_AMOUNT = USD(0.00001)
 
 
 class OmenAgentMarket(AgentMarket):
@@ -103,7 +97,6 @@ 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
 
@@ -152,25 +145,24 @@ class OmenAgentMarket(AgentMarket):
     def get_liquidity_in_wei(self, web3: Web3 | None = None) -> Wei:
         return self.get_contract().totalSupply(web3)
 
-    def get_liquidity_in_xdai(self, web3: Web3 | None = None) -> xDai:
-        return wei_to_xdai(self.get_liquidity_in_wei(web3))
+    def get_liquidity(self, web3: Web3 | None = None) -> Token:
+        return self.get_liquidity_in_wei(web3).as_token
 
-    def get_liquidity(self) -> TokenAmount:
-        return TokenAmount(
-            amount=self.get_liquidity_in_xdai(),
-            currency=self.currency,
-        )
+    def get_tiny_bet_amount(self) -> Token:
+        return self.get_in_token(OMEN_TINY_BET_AMOUNT)
 
-    @classmethod
-    def get_tiny_bet_amount(cls) -> BetAmount:
-        return BetAmount(amount=OMEN_TINY_BET_AMOUNT, currency=cls.currency)
+    def get_token_in_usd(self, x: Token) -> USD:
+        return get_token_in_usd(x, self.collateral_token_contract_address_checksummed)
+
+    def get_usd_in_token(self, x: USD) -> Token:
+        return get_usd_in_token(x, self.collateral_token_contract_address_checksummed)
 
     def liquidate_existing_positions(
         self,
         bet_outcome: bool,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
-        larger_than: float | None = None,
+        larger_than: OutcomeToken | None = None,
     ) -> None:
         """
         Liquidates all previously existing positions.
@@ -179,16 +171,14 @@ 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().amount
+            larger_than if larger_than is not None else self.get_liquidatable_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.items():
+            for position_outcome, token_amount in prev_position.amounts_ot.items():
                 position_outcome_bool = get_boolean_outcome(position_outcome)
                 if position_outcome_bool != bet_outcome:
                     self.sell_tokens(
@@ -202,7 +192,7 @@ class OmenAgentMarket(AgentMarket):
     def place_bet(
         self,
         outcome: bool,
-        amount: BetAmount,
+        amount: USD,
         auto_deposit: bool = True,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
@@ -211,12 +201,9 @@ 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_xdai,
+            amount=amount,
             market=self,
             binary_outcome=outcome,
             auto_deposit=auto_deposit,
@@ -226,7 +213,7 @@ class OmenAgentMarket(AgentMarket):
     def buy_tokens(
         self,
         outcome: bool,
-        amount: TokenAmount,
+        amount: USD,
         web3: Web3 | None = None,
         api_keys: APIKeys | None = None,
     ) -> str:
@@ -237,28 +224,37 @@ class OmenAgentMarket(AgentMarket):
             api_keys=api_keys,
         )
 
-    def calculate_sell_amount_in_collateral(
-        self, amount: TokenAmount, outcome: bool, web3: Web3 | None = None
-    ) -> xDai:
+    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)
+
         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 else self.no_index]
-        other_str = self.outcomes[self.no_index if outcome else self.yes_index]
+        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]
 
         collateral = calculate_sell_amount_in_collateral(
-            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)]),
+            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,
             fees=self.fees,
         )
-        return xDai(collateral)
+
+        return collateral
 
     def sell_tokens(
         self,
         outcome: bool,
-        amount: TokenAmount,
+        amount: USD | OutcomeToken,
         auto_withdraw: bool = True,
         api_keys: APIKeys | None = None,
         web3: Web3 | None = None,
@@ -267,16 +263,8 @@ class OmenAgentMarket(AgentMarket):
             raise ValueError(
                 f"Market {self.id} is not open for trading. Cannot sell tokens."
             )
-
-        # 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,
+            amount=amount,
             api_keys=api_keys if api_keys is not None else APIKeys(),
             market=self,
             binary_outcome=outcome,
@@ -324,7 +312,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=wei_type(0),
+            total_balance_bigger_than=OutcomeWei(0),
         )
         return len(user_positions) > 0
 
@@ -361,16 +349,16 @@ class OmenAgentMarket(AgentMarket):
             current_p_yes=model.current_p_yes,
             condition=model.condition,
             url=model.url,
-            volume=wei_to_xdai(model.collateralVolume),
+            volume=model.collateralVolume.as_token,
             close_time=model.close_time,
             fees=MarketFees(
                 bet_proportion=(
-                    float(wei_to_xdai(model.fee)) if model.fee is not None else 0.0
+                    model.fee.as_token.value if model.fee is not None else 0.0
                 ),
                 absolute=0,
             ),
             outcome_token_pool={
-                model.outcomes[i]: wei_to_xdai(Wei(model.outcomeTokenAmounts[i]))
+                model.outcomes[i]: model.outcomeTokenAmounts[i].as_outcome_token
                 for i in range(len(model.outcomes))
             },
         )
@@ -407,22 +395,24 @@ class OmenAgentMarket(AgentMarket):
         redeem_from_all_user_positions(api_keys)
 
     @staticmethod
-    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
+    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,
         )
+        return keeping_usd + native_usd
 
     @staticmethod
     def verify_operational_balance(api_keys: APIKeys) -> bool:
-        return get_total_balance(
-            api_keys.public_key,
+        return get_balances(
             # Use `public_key`, not `bet_from_address` because transaction costs are paid from the EOA wallet.
-            sum_wxdai=False,
-        ) > xdai_type(0.001)
+            api_keys.public_key,
+        ).xdai > xDai(0.001)
 
     def store_prediction(
         self, processed_market: ProcessedMarket | None, keys: APIKeys, agent_name: str
@@ -497,22 +487,19 @@ class OmenAgentMarket(AgentMarket):
 
     def get_token_balance(
         self, user_id: str, outcome: str, web3: Web3 | None = None
-    ) -> TokenAmount:
+    ) -> OutcomeToken:
         index_set = self.get_index_set(outcome)
         balances = get_conditional_tokens_balance_for_market(
             self, Web3.to_checksum_address(user_id), web3=web3
         )
-        return TokenAmount(
-            amount=wei_to_xdai(balances[index_set]),
-            currency=self.currency,
-        )
+        return balances[index_set].as_outcome_token
 
-    def get_position(self, user_id: str) -> Position | None:
+    def get_position(self, user_id: str) -> ExistingPosition | None:
         liquidatable_amount = self.get_liquidatable_amount()
         existing_positions = self.get_positions(
             user_id=user_id,
             liquid_only=True,
-            larger_than=liquidatable_amount.amount,
+            larger_than=liquidatable_amount,
         )
         existing_position = next(
             iter([i for i in existing_positions if i.market_id == self.id]), None
@@ -524,12 +511,12 @@ class OmenAgentMarket(AgentMarket):
         cls,
         user_id: str,
         liquid_only: bool = False,
-        larger_than: float = 0,
-    ) -> list[Position]:
+        larger_than: OutcomeToken = OutcomeToken(0),
+    ) -> t.Sequence[ExistingPosition]:
         sgh = OmenSubgraphHandler()
         omen_positions = sgh.get_user_positions(
             better_address=Web3.to_checksum_address(user_id),
-            total_balance_bigger_than=xdai_to_wei(xDai(larger_than)),
+            total_balance_bigger_than=larger_than.as_outcome_wei,
         )
 
         # Sort positions and corresponding markets by condition_id
@@ -563,142 +550,105 @@ class OmenAgentMarket(AgentMarket):
             if liquid_only and not market.can_be_traded():
                 continue
 
-            amounts: dict[OutcomeStr, TokenAmount] = {}
+            amounts_ot: dict[OutcomeStr, OutcomeToken] = {}
+
             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:
+                if outecome_str in amounts_ot:
                     raise ValueError(
-                        f"Outcome {outecome_str} already exists in {amounts=}"
+                        f"Outcome {outecome_str} already exists in {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}
+                amounts_ot[outecome_str] = omen_position.totalBalance.as_outcome_token
 
-        # 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()}"
+            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,
+                )
             )
 
-        # 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)
+        return positions
 
     @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: BetAmount, direction: bool
-    ) -> TokenAmount:
+        self, bet_amount: USD | Token, direction: bool
+    ) -> OutcomeToken:
         """
         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=bet_amount.amount,
+            investment_amount=self.get_in_token(bet_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 TokenAmount(amount=amount, currency=self.currency)
+        return amount
 
     def _get_buy_token_amount_from_smart_contract(
-        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)
-                ),
-            )
+        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 = float(wei_to_xdai(received_token_amount_wei))
-        return TokenAmount(amount=received_token_amount, currency=self.currency)
+        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)
+            ),
+        )
+        received_token_amount = received_token_amount_wei.as_outcome_token
+        return received_token_amount
 
-    def get_new_p_yes(self, bet_amount: BetAmount, direction: bool) -> Probability:
+    def get_new_p_yes(self, bet_amount: USD, 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)]
-        no_outcome_pool_size = outcome_token_pool[self.get_outcome_str_from_bool(False)]
+
+        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
 
         new_yes_outcome_pool_size = yes_outcome_pool_size + (
-            self.fees.get_bet_size_after_fees(bet_amount.amount)
+            self.fees.get_after_fees(bet_amount_in_tokens).value
         )
         new_no_outcome_pool_size = no_outcome_pool_size + (
-            self.fees.get_bet_size_after_fees(bet_amount.amount)
+            self.fees.get_after_fees(bet_amount_in_tokens).value
         )
 
-        received_token_amount = self.get_buy_token_amount(bet_amount, direction).amount
+        received_token_amount = self.get_buy_token_amount(bet_amount, direction).value
         if direction:
             new_yes_outcome_pool_size -= received_token_amount
         else:
@@ -737,11 +687,18 @@ def get_omen_user_url(address: ChecksumAddress) -> str:
 
 
 def pick_binary_market(
-    sort_by: SortBy = SortBy.CLOSING_SOONEST, filter_by: FilterBy = FilterBy.OPEN
+    sort_by: SortBy = SortBy.CLOSING_SOONEST,
+    filter_by: FilterBy = FilterBy.OPEN,
+    collateral_token_address_in: (
+        tuple[ChecksumAddress, ...] | None
+    ) = SAFE_COLLATERAL_TOKENS_ADDRESSES,
 ) -> OmenMarket:
     subgraph_handler = OmenSubgraphHandler()
     return subgraph_handler.get_omen_binary_markets_simple(
-        limit=1, sort_by=sort_by, filter_by=filter_by
+        limit=1,
+        sort_by=sort_by,
+        filter_by=filter_by,
+        collateral_token_address_in=collateral_token_address_in,
     )[0]
 
 
@@ -752,7 +709,7 @@ def pick_binary_market(
 )
 def omen_buy_outcome_tx(
     api_keys: APIKeys,
-    amount: xDai,
+    amount: USD | Token,
     market: OmenAgentMarket,
     outcome: str,
     auto_deposit: bool,
@@ -760,35 +717,32 @@ def omen_buy_outcome_tx(
     slippage: float = 0.01,
 ) -> str:
     """
-    Bets the given amount of xDai for the given outcome in the given market.
+    Bets the given amount 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)
 
-    # 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)
+    amount_token = market.get_in_token(amount)
+    amount_wei = amount_token.as_wei
 
     # 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_to_buy, outcome_index, web3=web3
+        amount_wei, outcome_index, web3=web3
     )
     # Allow small slippage.
-    expected_shares = remove_fraction(expected_shares, slippage)
+    expected_shares = expected_shares.without_fraction(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_to_buy,
+        amount_wei=amount_wei,
         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
         )
@@ -796,7 +750,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_to_buy,
+        amount_wei=amount_wei,
         outcome_index=outcome_index,
         min_outcome_tokens_to_buy=expected_shares,
         web3=web3,
@@ -807,7 +761,7 @@ def omen_buy_outcome_tx(
 
 def binary_omen_buy_outcome_tx(
     api_keys: APIKeys,
-    amount: xDai,
+    amount: USD | Token,
     market: OmenAgentMarket,
     binary_outcome: bool,
     auto_deposit: bool,
@@ -825,7 +779,7 @@ def binary_omen_buy_outcome_tx(
 
 def omen_sell_outcome_tx(
     api_keys: APIKeys,
-    amount: xDai,  # The xDai value of shares to sell.
+    amount: OutcomeToken | Token | USD,
     market: OmenAgentMarket,
     outcome: str,
     auto_withdraw: bool,
@@ -839,13 +793,18 @@ 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)
@@ -855,15 +814,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 buy.
+    # Get the index of the outcome we want to sell.
     outcome_index: int = market.get_outcome_index(outcome)
 
-    # Calculate the amount of shares we will sell for the given selling amount of xdai.
+    # Calculate the amount of shares we will sell for the given selling amount of collateral.
     max_outcome_tokens_to_sell = market_contract.calcSellAmount(
         amount_wei, outcome_index, web3=web3
     )
     # Allow small slippage.
-    max_outcome_tokens_to_sell = add_fraction(max_outcome_tokens_to_sell, slippage)
+    max_outcome_tokens_to_sell = max_outcome_tokens_to_sell.with_fraction(slippage)
 
     # Approve the market maker to move our (all) conditional tokens.
     conditional_token_contract.setApprovalForAll(
@@ -893,7 +852,7 @@ def omen_sell_outcome_tx(
 
 def binary_omen_sell_outcome_tx(
     api_keys: APIKeys,
-    amount: xDai,
+    amount: OutcomeToken | Token | USD,
     market: OmenAgentMarket,
     binary_outcome: bool,
     auto_withdraw: bool,
@@ -911,16 +870,16 @@ def binary_omen_sell_outcome_tx(
 
 def omen_create_market_tx(
     api_keys: APIKeys,
-    initial_funds: xDai,
+    initial_funds: USD | Token,
     question: str,
     closing_time: DatetimeUTC,
     category: str,
     language: str,
-    outcomes: list[str],
+    outcomes: t.Sequence[str],
     auto_deposit: bool,
     finalization_timeout: timedelta = REALITY_DEFAULT_FINALIZATION_TIMEOUT,
     fee_perc: float = OMEN_DEFAULT_MARKET_FEE_PERC,
-    distribution_hint: list[OmenOutcomeToken] | None = None,
+    distribution_hint: list[OutcomeWei] | None = None,
     collateral_token_address: ChecksumAddress = WrappedxDaiContract().address,
     arbitrator: Arbitrator = Arbitrator.KLEROS_31_JURORS_WITH_APPEAL,
     web3: Web3 | None = None,
@@ -931,7 +890,12 @@ def omen_create_market_tx(
     web3 = (
         web3 or OmenFixedProductMarketMakerFactoryContract.get_web3()
     )  # Default to Gnosis web3.
-    initial_funds_wei = xdai_to_wei(initial_funds)
+    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
 
     realitio_contract = OmenRealitioContract()
     conditional_token_contract = OmenConditionalTokenContract()
@@ -955,9 +919,9 @@ def omen_create_market_tx(
 
     if auto_deposit:
         auto_deposit_collateral_token(
-            collateral_token_contract=collateral_token_contract,
+            collateral_token_contract,
+            initial_funds_in_collateral_wei,
             api_keys=api_keys,
-            amount_wei=initial_funds_wei,
             web3=web3,
         )
 
@@ -991,21 +955,16 @@ 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_shares,
+        amount_wei=initial_funds_in_collateral_wei,
         web3=web3,
     )
 
     # Create the market.
-    fee = xdai_to_wei(xdai_type(fee_perc))
+    fee = Token(fee_perc).as_wei
     (
         market_event,
         funding_event,
@@ -1015,7 +974,7 @@ def omen_create_market_tx(
         condition_id=condition_id,
         fee=fee,
         distribution_hint=distribution_hint,
-        initial_funds_wei=initial_funds_in_shares,
+        initial_funds_wei=initial_funds_in_collateral_wei,
         collateral_token_address=collateral_token_contract.address,
         web3=web3,
     )
@@ -1032,7 +991,7 @@ def omen_create_market_tx(
         condition_id=condition_id,
         question_event=question_event,
         condition_event=cond_event,
-        initial_funds=initial_funds_wei,
+        initial_funds=initial_funds_in_collateral_wei,
         fee=fee,
         distribution_hint=distribution_hint,
     )
@@ -1041,27 +1000,28 @@ def omen_create_market_tx(
 def omen_fund_market_tx(
     api_keys: APIKeys,
     market: OmenAgentMarket,
-    funds: Wei,
+    funds: USD | Token,
     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=amount_to_fund,
+        amount_wei=funds_in_collateral_wei,
         web3=web3,
     )
 
     if auto_deposit:
-        # 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)
+        auto_deposit_collateral_token(
+            collateral_token_contract, funds_in_collateral_wei, api_keys, web3
+        )
 
-    market_contract.addFunding(api_keys, amount_to_fund, web3=web3)
+    market_contract.addFunding(api_keys, funds_in_collateral_wei, web3=web3)
 
 
 def omen_redeem_full_position_tx(
@@ -1094,7 +1054,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())
+    amount_wei = sum(amount_per_index.values(), start=OutcomeWei.zero())
     if amount_wei == 0:
         logger.debug("No balance to claim. Exiting.")
         return
@@ -1112,10 +1072,10 @@ def omen_redeem_full_position_tx(
         web3=web3,
     )
     new_balance = collateral_token_contract.balanceOf(from_address, web3=web3)
-    balance_diff = wei_type(new_balance - original_balance)
+    balance_diff = new_balance - original_balance
 
     logger.info(
-        f"Redeemed {wei_to_xdai(balance_diff)} {collateral_token_contract.symbol_cached(web3=web3)} from market {market.question=} ({market.url})."
+        f"Redeemed {balance_diff.as_token} {collateral_token_contract.symbol_cached(web3=web3)} from market {market.question=} ({market.url})."
     )
 
     if auto_withdraw:
@@ -1131,12 +1091,12 @@ def get_conditional_tokens_balance_for_market(
     market: OmenAgentMarket,
     from_address: ChecksumAddress,
     web3: Web3 | None = None,
-) -> dict[int, Wei]:
+) -> dict[int, OutcomeWei]:
     """
     We derive the withdrawable balance from the ConditionalTokens contract through CollectionId -> PositionId (which
     also serves as tokenId) -> TokenBalances.
     """
-    balance_per_index_set: dict[int, Wei] = {}
+    balance_per_index_set: dict[int, OutcomeWei] = {}
     conditional_token_contract = OmenConditionalTokenContract()
     parent_collection_id = build_parent_collection_id()
 
@@ -1153,7 +1113,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] = wei_type(balance_for_position)
+        balance_per_index_set[index_set] = balance_for_position
 
     return balance_per_index_set
 
@@ -1214,11 +1174,11 @@ def omen_remove_fund_market_tx(
     )
 
     new_balance = market_collateral_token_contract.balanceOf(from_address, web3=web3)
-    balance_diff = wei_type(new_balance - original_balance)
+    balance_diff = new_balance - original_balance
 
     logger.debug(f"Result from merge positions {result}")
     logger.info(
-        f"Withdrawn {wei_to_xdai(balance_diff)} {market_collateral_token_contract.symbol_cached(web3=web3)} from liquidity at {market.url=}."
+        f"Withdrawn {balance_diff.as_token} {market_collateral_token_contract.symbol_cached(web3=web3)} from liquidity at {market.url=}."
     )
 
     if auto_withdraw:
@@ -1244,7 +1204,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=wei_type(0),
+        total_balance_bigger_than=OutcomeWei(0),
     )
 
     for index, user_position in enumerate(user_positions):
@@ -1272,10 +1232,10 @@ def redeem_from_all_user_positions(
             web3=web3,
         )
         new_balance = collateral_token_contract.balanceOf(public_key, web3=web3)
-        balance_diff = wei_type(new_balance - original_balance)
+        balance_diff = new_balance - original_balance
 
         logger.info(
-            f"Redeemed {wei_to_xdai(balance_diff)} {collateral_token_contract.symbol_cached(web3=web3)} from position {user_position.id=}."
+            f"Redeemed {balance_diff.as_token} {collateral_token_contract.symbol_cached(web3=web3)} from position {user_position.id=}."
         )
 
         if auto_withdraw:
@@ -1346,23 +1306,21 @@ def send_keeping_token_to_eoa_xdai(
         )
         return
 
-    need_to_withdraw = xDai(
-        (min_required_balance - current_balances_eoa.xdai) * multiplier
-    )
-    need_to_withdraw_wei = xdai_to_wei(need_to_withdraw)
+    need_to_withdraw = (min_required_balance - current_balances_eoa.xdai) * multiplier
+    need_to_withdraw_wei = need_to_withdraw.as_xdai_wei
 
-    if current_balances_eoa.wxdai >= need_to_withdraw:
+    if current_balances_eoa.wxdai >= need_to_withdraw.as_token:
         # 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,
+            amount_wei=need_to_withdraw_wei.as_wei,
             web3=web3,
         )
 
-    elif current_balances_betting.wxdai >= need_to_withdraw:
+    elif current_balances_betting.wxdai >= need_to_withdraw.as_token:
         # If Safe has enough of wxDai:
         # First send them to EOA's address.
         logger.info(
@@ -1372,7 +1330,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,
+            amount_wei=need_to_withdraw_wei.as_wei,
             web3=web3,
         )
         # And then simply withdraw it.
@@ -1381,7 +1339,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,
+            amount_wei=need_to_withdraw_wei.as_wei,
             web3=web3,
         )
 
@@ -1392,30 +1350,37 @@ def send_keeping_token_to_eoa_xdai(
 
 
 def get_buy_outcome_token_amount(
-    investment_amount: float,
+    investment_amount: Token,
     buy_direction: bool,
-    yes_outcome_pool_size: float,
-    no_outcome_pool_size: float,
+    yes_outcome_pool_size: OutcomeToken,
+    no_outcome_pool_size: OutcomeToken,
     fees: MarketFees,
-) -> float:
+) -> OutcomeToken:
     """
     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_bet_size_after_fees(investment_amount)
+    investment_amount_minus_fees = fees.get_after_fees(investment_amount)
+    investment_amount_minus_fees_as_ot = OutcomeToken(
+        investment_amount_minus_fees.value
+    )
     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
-    ending_outcome_balance = buy_token_pool_balance * pool_balance / denominator
+    denominator = pool_balance + investment_amount_minus_fees_as_ot
+    ending_outcome_balance = OutcomeToken(
+        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 - ending_outcome_balance
+        buy_token_pool_balance
+        + investment_amount_minus_fees_as_ot
+        - ending_outcome_balance
     )
     return result