Added WithdrawFromTreasury function (#643)

* Added WithdrawFromTreasury function

* Updated PMAT

* Fixed mypy

* Small fixes after merge

prediction_market_agent/agents/microchain_agent/nft_functions.py

@@ -2,6 +2,7 @@
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.contract import (
     ContractOwnableERC721OnGnosisChain,
+    SimpleTreasuryContract,
 )
 from web3 import Web3
 
@@ -51,7 +52,7 @@ def __call__(
         contract = ContractOwnableERC721OnGnosisChain(
             address=Web3.to_checksum_address(nft_address)
         )
-        owner_address = contract.ownerOf(token_id)
+        owner_address: str = contract.owner_of(token_id)
         return owner_address
 
 
@@ -90,8 +91,29 @@ def __call__(
         return "Token transferred successfully."
 
 
+class WithdrawFromTreasury(Function):
+    @property
+    def description(self) -> str:
+        required_balance_nft_tokens = SimpleTreasuryContract().required_nft_balance()
+        return f"Transfers the entire balance of the treasury to the caller. For the function to succeed, the caller must own {required_balance_nft_tokens} NFT tokens."
+
+    @property
+    def example_args(self) -> list[str]:
+        return []
+
+    def __call__(self) -> str:
+        keys = MicrochainAgentKeys()
+        treasury_contract = SimpleTreasuryContract()
+        logger.info(
+            f"Withdrawing from the treasury using sender {keys.bet_from_address}"
+        )
+        treasury_contract.withdraw(api_keys=keys)
+        return "Treasury successfully emptied."
+
+
 NFT_FUNCTIONS: list[type[Function]] = [
     BalanceOfNFT,
     OwnerOfNFT,
     SafeTransferFromNFT,
+    WithdrawFromTreasury,
 ]

prediction_market_agent/agents/microchain_agent/nft_treasury_game/app_nft_treasury_game.py

@@ -25,7 +25,7 @@
 from prediction_market_agent.agents.microchain_agent.nft_functions import BalanceOfNFT
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.constants_nft_treasury_game import (
     NFT_TOKEN_FACTORY,
-    TREASURY_SAFE_ADDRESS,
+    TREASURY_ADDRESS,
 )
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.contracts_nft_treasury_game import (
     get_nft_token_factory_max_supply,
@@ -310,7 +310,7 @@ def show_about_agent_part(nft_agent: type[DeployableAgentNFTGameAbstract]) -> No
 
 @st.fragment(run_every=timedelta(seconds=10))
 def show_treasury_part() -> None:
-    treasury_xdai_balance = get_balances(TREASURY_SAFE_ADDRESS).xdai
+    treasury_xdai_balance = get_balances(TREASURY_ADDRESS).xdai
     st.markdown(
         f"""### Treasury
 Currently holds <span style='font-size: 1.1em;'><strong>{treasury_xdai_balance:.2f} xDAI</strong></span>. There are {get_nft_token_factory_max_supply()} NFT keys.""",

prediction_market_agent/agents/microchain_agent/nft_treasury_game/constants_nft_treasury_game.py

@@ -1,7 +1,7 @@
 from web3 import Web3
 
-TREASURY_SAFE_ADDRESS = Web3.to_checksum_address(
-    "0xd1A54FD7a200C2ca76B6D06437795d660d37FE28"
+TREASURY_ADDRESS = Web3.to_checksum_address(
+    "0x624ad0db52e6b18afb4d36b8e79d0c2a74f3fc8a"
 )
 NFT_TOKEN_FACTORY = Web3.to_checksum_address(
     "0x0D7C0Bd4169D090038c6F41CFd066958fe7619D0"

prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py

@@ -25,7 +25,7 @@ def token_ids_owned_by(
         return [
             token_id
             for token_id in token_ids
-            if self.ownerOf(tokenId=token_id, web3=web3) == owner
+            if self.owner_of(token_id=token_id, web3=web3) == owner
         ]
 
 

prediction_market_agent/agents/microchain_agent/nft_treasury_game/deploy_nft_treasury_game.py

@@ -1,12 +1,9 @@
 import time
 
-from eth_typing import URI
 from microchain.functions import Reasoning
-from prediction_market_agent_tooling.config import RPCConfig
 from prediction_market_agent_tooling.gtypes import ChecksumAddress
 from prediction_market_agent_tooling.loggers import logger
-from safe_eth.eth import EthereumClient
-from safe_eth.safe.safe import Safe, SafeV141
+from prediction_market_agent_tooling.tools.contract import SimpleTreasuryContract
 from web3 import Web3
 
 from prediction_market_agent.agents.identifiers import AgentIdentifier
@@ -26,7 +23,7 @@
 )
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.constants_nft_treasury_game import (
     NFT_TOKEN_FACTORY,
-    TREASURY_SAFE_ADDRESS,
+    TREASURY_ADDRESS,
 )
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.contracts_nft_treasury_game import (
     get_nft_token_factory_max_supply,
@@ -61,20 +58,15 @@ class DeployableAgentNFTGameAbstract(DeployableMicrochainAgentAbstract):
     # Game status
     game_finished_detected: bool = False
 
-    @classmethod
-    def build_treasury_safe(cls) -> Safe:
-        client = EthereumClient(URI(RPCConfig().gnosis_rpc_url))
-        return SafeV141(TREASURY_SAFE_ADDRESS, client)
-
     @classmethod
     def retrieve_treasury_thresold(cls) -> int:
-        safe = cls.build_treasury_safe()
-        return safe.retrieve_threshold()
+        return SimpleTreasuryContract().required_nft_balance()
 
     @classmethod
-    def retrieve_treasury_owners(cls) -> list[ChecksumAddress]:
-        safe = cls.build_treasury_safe()
-        return [Web3.to_checksum_address(o) for o in safe.retrieve_owners()]
+    def retrieve_total_number_of_keys(cls) -> int:
+        # We could iteratively call `owner_of` for a range of token_ids, thus finding out the max supply. However,
+        # in the current implementation, no new tokens can be created and max_supply = 5, hence hardcoding it here.
+        return 5
 
     @classmethod
     def get_description(cls) -> str:
@@ -333,7 +325,7 @@ def nft_treasury_game_base_prompt(wallet_address: ChecksumAddress) -> str:
     - The current amount of xDai in the treasury
     - The total number of keys in circulation
     - The distribution of keys among participants
-- Address of the treasury, wallet holding the treasury's xDai, is {TREASURY_SAFE_ADDRESS}.
+- Address of the treasury, wallet holding the treasury's xDai, is {TREASURY_ADDRESS}.
 - Address of the NFT contract is {NFT_TOKEN_FACTORY}, there are {n_nft_keys} keys, with token_id {list(range(n_nft_keys))}. 
   - You can own multiple NFT keys. 
   - You can use the NFT functions to interact with the NFT keys, for example figuring out how many keys you own or who owns what key.
@@ -354,7 +346,7 @@ def nft_treasury_game_buyer_prompt() -> str:
     return f"""You participate in the NFT Treasury game as a key buyer.
 
 [OBJECTIVE]
-- Your goal is to acquire {DeployableAgentNFTGameAbstract.retrieve_treasury_thresold()} out of {len(DeployableAgentNFTGameAbstract.retrieve_treasury_owners())} NFT keys to unlock the treasury
+- Your goal is to acquire {DeployableAgentNFTGameAbstract.retrieve_treasury_thresold()} out of {DeployableAgentNFTGameAbstract.retrieve_total_number_of_keys()} NFT keys to unlock the treasury
 - The total xDai spent on acquiring these keys must be less than the treasury's value to ensure a profitable outcome when claiming the treasury.
 
 [KEY ACQUISITION STRATEGY]

prediction_market_agent/agents/microchain_agent/nft_treasury_game/tools_nft_treasury_game.py

@@ -4,7 +4,7 @@
 from prediction_market_agent_tooling.tools.balances import get_balances
 
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.constants_nft_treasury_game import (
-    TREASURY_SAFE_ADDRESS,
+    TREASURY_ADDRESS,
 )
 
 
@@ -14,7 +14,7 @@ class NFTGameStatus(str, Enum):
 
 
 def get_nft_game_status() -> NFTGameStatus:
-    treasury_balance = get_balances(TREASURY_SAFE_ADDRESS)
+    treasury_balance = get_balances(TREASURY_ADDRESS)
 
     if treasury_balance.total == xDai(0):
         return NFTGameStatus.finished

pyproject.toml

@@ -33,7 +33,7 @@ poetry = "^1.7.1"
 poetry-plugin-export = "^1.6.0"
 functions-framework = "^3.5.0"
 cron-validator = "^1.0.8"
-prediction-market-agent-tooling = { version = "^0.57.11", extras = ["langchain", "google"] }
+prediction-market-agent-tooling = "0.57.12"
 pydantic-settings = "^2.1.0"
 autoflake = "^2.2.1"
 isort = "^5.13.2"