prediction-market-agent-tooling 0.59.2__py3-none-any.whl → 0.60.0.dev408__py3-none-any.whl

prediction_market_agent_tooling/markets/omen/data_models.py

@@ -23,6 +23,11 @@ from prediction_market_agent_tooling.markets.data_models import (
     Resolution,
     ResolvedBet,
 )
+from prediction_market_agent_tooling.tools.contract import (
+    ContractERC20OnGnosisChain,
+    init_collateral_token_contract,
+    to_gnosis_chain_contract,
+)
 from prediction_market_agent_tooling.tools.utils import (
     BPS_CONSTANT,
     DatetimeUTC,
@@ -168,6 +173,16 @@ class OmenPosition(BaseModel):
     def collateral_token_contract_address_checksummed(self) -> ChecksumAddress:
         return Web3.to_checksum_address(self.collateralTokenAddress)
 
+    def get_collateral_token_contract(
+        self, web3: Web3 | None
+    ) -> ContractERC20OnGnosisChain:
+        web3 = web3 or ContractERC20OnGnosisChain.get_web3()
+        return to_gnosis_chain_contract(
+            init_collateral_token_contract(
+                self.collateral_token_contract_address_checksummed, web3
+            )
+        )
+
 
 class OmenUserPosition(BaseModel):
     id: HexBytes

prediction_market_agent_tooling/markets/omen/omen.py

@@ -80,6 +80,7 @@ from prediction_market_agent_tooling.tools.tokens.auto_deposit import (
 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.utils import (
     DatetimeUTC,
     calculate_sell_amount_in_collateral,
@@ -190,11 +191,10 @@ class OmenAgentMarket(AgentMarket):
             for position_outcome, token_amount in prev_position.amounts.items():
                 position_outcome_bool = get_boolean_outcome(position_outcome)
                 if position_outcome_bool != bet_outcome:
-                    # We keep it as collateral since we want to place a bet immediately after this function.
                     self.sell_tokens(
                         outcome=position_outcome_bool,
                         amount=token_amount,
-                        auto_withdraw=False,
+                        auto_withdraw=True,
                         web3=web3,
                         api_keys=api_keys,
                     )
@@ -259,7 +259,7 @@ class OmenAgentMarket(AgentMarket):
         self,
         outcome: bool,
         amount: TokenAmount,
-        auto_withdraw: bool = False,
+        auto_withdraw: bool = True,
         api_keys: APIKeys | None = None,
         web3: Web3 | None = None,
     ) -> str:
@@ -408,8 +408,8 @@ class OmenAgentMarket(AgentMarket):
 
     @staticmethod
     def get_trade_balance(api_keys: APIKeys, web3: Web3 | None = None) -> xDai:
-        return get_total_balance(
-            address=api_keys.bet_from_address, web3=web3, sum_xdai=True, sum_wxdai=True
+        return wei_to_xdai(
+            KEEPING_ERC20_TOKEN.balanceOf(api_keys.bet_from_address, web3=web3)
         )
 
     @staticmethod
@@ -1060,6 +1060,7 @@ def omen_fund_market_tx(
 def omen_redeem_full_position_tx(
     api_keys: APIKeys,
     market: OmenAgentMarket,
+    auto_withdraw: bool = True,
     web3: Web3 | None = None,
 ) -> None:
     """
@@ -1071,6 +1072,7 @@ def omen_redeem_full_position_tx(
 
     market_contract: OmenFixedProductMarketMakerContract = market.get_contract()
     conditional_token_contract = OmenConditionalTokenContract()
+    collateral_token_contract = market_contract.get_collateral_token_contract(web3)
 
     # Verify, that markets uses conditional tokens that we expect.
     if market_contract.conditionalTokens() != conditional_token_contract.address:
@@ -1094,6 +1096,7 @@ def omen_redeem_full_position_tx(
         logger.debug("Market not yet resolved, not possible to claim")
         return
 
+    original_balance = collateral_token_contract.balanceOf(from_address, web3=web3)
     conditional_token_contract.redeemPositions(
         api_keys=api_keys,
         collateral_token_address=market.collateral_token_contract_address_checksummed,
@@ -1101,6 +1104,20 @@ def omen_redeem_full_position_tx(
         index_sets=market.condition.index_sets,
         web3=web3,
     )
+    new_balance = collateral_token_contract.balanceOf(from_address, web3=web3)
+    balance_diff = wei_type(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})."
+    )
+
+    if auto_withdraw:
+        auto_withdraw_collateral_token(
+            collateral_token_contract=collateral_token_contract,
+            amount_wei=balance_diff,
+            api_keys=api_keys,
+            web3=web3,
+        )
 
 
 def get_conditional_tokens_balance_for_market(
@@ -1139,13 +1156,14 @@ def omen_remove_fund_market_tx(
     market: OmenAgentMarket,
     shares: Wei | None,
     web3: Web3 | None = None,
+    auto_withdraw: bool = True,
 ) -> None:
     """
     Removes funding from a given OmenMarket (moving the funds from the OmenMarket to the
     ConditionalTokens contract), and finally calls the `mergePositions` method which transfers collateralToken from the ConditionalTokens contract to the address corresponding to `from_private_key`.
 
     Warning: Liquidity removal works on the principle of getting market's shares, not the collateral token itself.
-    After we remove funding, using the `mergePositions` we get `min(shares per index)` of wxDai back, but the remaining shares can be converted back only after the market is resolved.
+    After we remove funding, using the `mergePositions` we get `min(shares per index)` of collateral token back, but the remaining shares can be converted back only after the market is resolved.
     That can be done using the `redeem_from_all_user_positions` function below.
     """
     from_address = api_keys.bet_from_address
@@ -1189,16 +1207,26 @@ 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)
 
     logger.debug(f"Result from merge positions {result}")
     logger.info(
-        f"Withdrawn {new_balance - original_balance} {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:
+        auto_withdraw_collateral_token(
+            collateral_token_contract=market_collateral_token_contract,
+            amount_wei=balance_diff,
+            api_keys=api_keys,
+            web3=web3,
+        )
+
 
 def redeem_from_all_user_positions(
     api_keys: APIKeys,
     web3: Web3 | None = None,
+    auto_withdraw: bool = True,
 ) -> None:
     """
     Redeems from all user positions where the user didn't redeem yet.
@@ -1224,8 +1252,11 @@ def redeem_from_all_user_positions(
         logger.info(
             f"[{index + 1} / {len(user_positions)}] Processing redeem from {user_position.id=}."
         )
+        collateral_token_contract = (
+            user_position.position.get_collateral_token_contract(web3=web3)
+        )
 
-        original_balances = get_balances(public_key, web3)
+        original_balance = collateral_token_contract.balanceOf(public_key, web3=web3)
         conditional_token_contract.redeemPositions(
             api_keys=api_keys,
             collateral_token_address=user_position.position.collateral_token_contract_address_checksummed,
@@ -1233,12 +1264,21 @@ def redeem_from_all_user_positions(
             index_sets=user_position.position.indexSets,
             web3=web3,
         )
-        new_balances = get_balances(public_key, web3)
+        new_balance = collateral_token_contract.balanceOf(public_key, web3=web3)
+        balance_diff = wei_type(new_balance - original_balance)
 
         logger.info(
-            f"Redeemed {new_balances.wxdai - original_balances.wxdai} wxDai 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:
+            auto_withdraw_collateral_token(
+                collateral_token_contract=collateral_token_contract,
+                amount_wei=balance_diff,
+                api_keys=api_keys,
+                web3=web3,
+            )
+
 
 def get_binary_market_p_yes_history(market: OmenAgentMarket) -> list[Probability]:
     history: list[Probability] = []

prediction_market_agent_tooling/tools/omen/sell_positions.py

@@ -13,7 +13,7 @@ from prediction_market_agent_tooling.tools.utils import utcnow
 def sell_all(
     api_keys: APIKeys,
     closing_later_than_days: int,
-    auto_withdraw: bool = False,
+    auto_withdraw: bool = True,
 ) -> None:
     """
     Helper function to sell all existing outcomes on Omen that would resolve later than in X days.

prediction_market_agent_tooling/tools/tokens/auto_withdraw.py

@@ -2,6 +2,7 @@ from web3 import Web3
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.gtypes import Wei
+from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.contract import (
     ContractERC20BaseClass,
     ContractERC4626BaseClass,
@@ -12,7 +13,10 @@ from prediction_market_agent_tooling.tools.cow.cow_order import (
 )
 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
+from prediction_market_agent_tooling.tools.web3_utils import (
+    remove_fraction,
+    wei_to_xdai,
+)
 
 
 def auto_withdraw_collateral_token(
@@ -28,8 +32,17 @@ def auto_withdraw_collateral_token(
         slippage,
     )
 
+    if not amount_wei:
+        logger.warning(
+            f"Amount to withdraw is zero, skipping withdrawal of {collateral_token_contract.symbol_cached(web3)}."
+        )
+        return
+
     if collateral_token_contract.address == KEEPING_ERC20_TOKEN.address:
         # Do nothing, as this is the token we want to keep.
+        logger.info(
+            f"Collateral token {collateral_token_contract.symbol_cached(web3)} is the same as KEEPING_ERC20_TOKEN. Not withdrawing."
+        )
         return
     elif (
         isinstance(collateral_token_contract, ContractERC4626BaseClass)
@@ -37,12 +50,18 @@ def auto_withdraw_collateral_token(
         == KEEPING_ERC20_TOKEN.address
     ):
         # 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)}"
+        )
         collateral_token_contract.withdraw(
             api_keys,
             amount_wei,
             web3=web3,
         )
     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)}"
+        )
         # 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(

{prediction_market_agent_tooling-0.59.2.dist-info → prediction_market_agent_tooling-0.60.0.dev408.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: prediction-market-agent-tooling
-Version: 0.59.2
+Version: 0.60.0.dev408
 Summary: Tools to benchmark, deploy and monitor prediction market agents.
 Author: Gnosis
 Requires-Python: >=3.10,<3.13

{prediction_market_agent_tooling-0.59.2.dist-info → prediction_market_agent_tooling-0.60.0.dev408.dist-info}/RECORD

@@ -50,8 +50,8 @@ prediction_market_agent_tooling/markets/metaculus/api.py,sha256=4TRPGytQQbSdf42D
 prediction_market_agent_tooling/markets/metaculus/data_models.py,sha256=FaBCTPPezXbBwZ9p791CiVgQ4vB696xnMbz9XVXmiVI,3267
 prediction_market_agent_tooling/markets/metaculus/metaculus.py,sha256=86TIx6cavEWc8Cv4KpZxSvwiSw9oFybXE3YB49pg-CA,4377
 prediction_market_agent_tooling/markets/omen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-prediction_market_agent_tooling/markets/omen/data_models.py,sha256=uT8ILKrg2g4jGodPxtolPErk25buNzMYndb01ZL2dYE,28421
-prediction_market_agent_tooling/markets/omen/omen.py,sha256=mtkLJR1w3QRfyW6aY_Ey-3ITL2BPHrisvPY0tmB4wGU,49864
+prediction_market_agent_tooling/markets/omen/data_models.py,sha256=sfaOpNk6oFIzxYQzs9EehqAT_19IxYJy9pns-UTepOc,28934
+prediction_market_agent_tooling/markets/omen/omen.py,sha256=7XOcjzuzX4vJUVr-M2yp1Diq8MaURFsC_8syhdpZbGg,51460
 prediction_market_agent_tooling/markets/omen/omen_constants.py,sha256=D9oflYKafLQiHYtB5sScMHqmXyzM8JP8J0yATmc4SQQ,233
 prediction_market_agent_tooling/markets/omen/omen_contracts.py,sha256=EXqBlVivbmW8aBQ65O09X2xkyesHAop49GUl1tUffWA,28648
 prediction_market_agent_tooling/markets/omen/omen_resolving.py,sha256=E1BVDvQd1qYtCxmfC94kJtGkmQqpGPHL3zTkcs5wW6M,9697
@@ -101,7 +101,7 @@ prediction_market_agent_tooling/tools/is_predictable.py,sha256=qVd6zqay2Dg2fyeAu
 prediction_market_agent_tooling/tools/langfuse_.py,sha256=jI_4ROxqo41CCnWGS1vN_AeDVhRzLMaQLxH3kxDu3L8,1153
 prediction_market_agent_tooling/tools/langfuse_client_utils.py,sha256=83T31s-YbsnNBLyYCjmBI2BBKUEqJUuYFa0uCdkoqy8,5901
 prediction_market_agent_tooling/tools/omen/reality_accuracy.py,sha256=M1SF7iSW1gVlQSTskdVFTn09uPLST23YeipVIWj54io,2236
-prediction_market_agent_tooling/tools/omen/sell_positions.py,sha256=JIfTMKF0bHyQ9qXNfyr5ul_eu4hBuydlRVGJET8Uw90,2308
+prediction_market_agent_tooling/tools/omen/sell_positions.py,sha256=hZCxXpcACO95DyiZ5oLFp982N0erZg4wccdSUKTgRlA,2307
 prediction_market_agent_tooling/tools/parallelism.py,sha256=6Gou0hbjtMZrYvxjTDFUDZuxmE2nqZVbb6hkg1hF82A,1022
 prediction_market_agent_tooling/tools/relevant_news_analysis/data_models.py,sha256=95l84aztFaxcRLLcRQ46yKJbIlOEuDAbIGLouyliDzA,1316
 prediction_market_agent_tooling/tools/relevant_news_analysis/relevant_news_analysis.py,sha256=CddJem7tk15NAudJDwmuL8znTycbR-YI8kTNtd3LzG8,5474
@@ -112,13 +112,13 @@ prediction_market_agent_tooling/tools/streamlit_user_login.py,sha256=NXEqfjT9Lc9
 prediction_market_agent_tooling/tools/tavily/tavily_models.py,sha256=5ldQs1pZe6uJ5eDAuP4OLpzmcqYShlIV67kttNFvGS0,342
 prediction_market_agent_tooling/tools/tavily/tavily_search.py,sha256=Kw2mXNkMTYTEe1MBSTqhQmLoeXtgb6CkmHlcAJvhtqE,3809
 prediction_market_agent_tooling/tools/tokens/auto_deposit.py,sha256=o8_ERfPL-ps9FLvH5vgdiSRJQ4dZONJw9KK9sHgeP2I,6390
-prediction_market_agent_tooling/tools/tokens/auto_withdraw.py,sha256=B02maQkl3wNptWFtZdnDosICJE5BfnLUVqxp5uJIaPA,2353
+prediction_market_agent_tooling/tools/tokens/auto_withdraw.py,sha256=gOvfZlrw6hdpYquGKLjMw-YpjAvWHnyGa-ke0Z6GIzU,3150
 prediction_market_agent_tooling/tools/tokens/main_token.py,sha256=5iHO7-iehSlXuue6ocVrr4IsklVjm7QHIwln4BsebJA,573
 prediction_market_agent_tooling/tools/transaction_cache.py,sha256=K5YKNL2_tR10Iw2TD9fuP-CTGpBbZtNdgbd0B_R7pjg,1814
 prediction_market_agent_tooling/tools/utils.py,sha256=jLG4nbEoIzzJiZ4RgMx4Q969Zdl0p0s63p8uET_0Fuw,6440
 prediction_market_agent_tooling/tools/web3_utils.py,sha256=2PXZfGRrDVZD60agVpBN4JkOF0YsNBXgTEH1y-V71uQ,12723
-prediction_market_agent_tooling-0.59.2.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
-prediction_market_agent_tooling-0.59.2.dist-info/METADATA,sha256=MqNzGgXLVbHhIWejKukTi6hRsjK0nr5TRbax5iiyLn8,8629
-prediction_market_agent_tooling-0.59.2.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-prediction_market_agent_tooling-0.59.2.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
-prediction_market_agent_tooling-0.59.2.dist-info/RECORD,,
+prediction_market_agent_tooling-0.60.0.dev408.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
+prediction_market_agent_tooling-0.60.0.dev408.dist-info/METADATA,sha256=_g5AhhMSsCOT5zVXMNwOPIVxIZGpPhP8jTsd1CaTPpY,8636
+prediction_market_agent_tooling-0.60.0.dev408.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+prediction_market_agent_tooling-0.60.0.dev408.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
+prediction_market_agent_tooling-0.60.0.dev408.dist-info/RECORD,,