From c38e310991351cf641dedb5be646e07e7d385d76 Mon Sep 17 00:00:00 2001
From: Peter Jung <peter@jung.ninja>
Date: Fri, 17 Jan 2025 15:41:07 +0000
Subject: [PATCH] Fix NFT contract usage

---
 .../app_nft_treasury_game.py                  |  5 +----
 .../contracts_nft_treasury_game.py            | 22 -------------------
 .../deploy_nft_treasury_game.py               |  5 +----
 3 files changed, 2 insertions(+), 30 deletions(-)

diff --git a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/app_nft_treasury_game.py b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/app_nft_treasury_game.py
index 6dcf02ae..447ce048 100644
--- a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/app_nft_treasury_game.py
+++ b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/app_nft_treasury_game.py
@@ -27,9 +27,6 @@
     NFT_TOKEN_FACTORY,
     TREASURY_ADDRESS,
 )
-from prediction_market_agent.agents.microchain_agent.nft_treasury_game.contracts_nft_treasury_game import (
-    get_nft_token_factory_max_supply,
-)
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.deploy_nft_treasury_game import (
     DEPLOYED_NFT_AGENTS,
     DeployableAgentNFTGameAbstract,
@@ -313,7 +310,7 @@ def show_treasury_part() -> None:
     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.""",
+Currently holds <span style='font-size: 1.1em;'><strong>{treasury_xdai_balance:.2f} xDAI</strong></span>. There are {DeployableAgentNFTGameAbstract.retrieve_total_number_of_keys()} NFT keys.""",
         unsafe_allow_html=True,
     )
 
diff --git a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py
index d39ea890..a8cc6d8a 100644
--- a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py
+++ b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py
@@ -1,10 +1,7 @@
-from functools import cache
-
 from prediction_market_agent_tooling.gtypes import ChecksumAddress
 from prediction_market_agent_tooling.tools.contract import (
     ContractOwnableERC721OnGnosisChain,
 )
-from web3 import Web3
 
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.constants_nft_treasury_game import (
     NFT_TOKEN_FACTORY,
@@ -13,22 +10,3 @@
 
 class ContractNFTFactoryOnGnosisChain(ContractOwnableERC721OnGnosisChain):
     address: ChecksumAddress = NFT_TOKEN_FACTORY
-
-    def max_supply(self, web3: Web3 | None = None) -> int:
-        n_tokens: int = self.call("MAX_SUPPLY", web3=web3)
-        return n_tokens
-
-    def token_ids_owned_by(
-        self, owner: ChecksumAddress, web3: Web3 | None = None
-    ) -> list[int]:
-        token_ids = list(range(self.max_supply(web3=web3)))
-        return [
-            token_id
-            for token_id in token_ids
-            if self.owner_of(token_id=token_id, web3=web3) == owner
-        ]
-
-
-@cache
-def get_nft_token_factory_max_supply() -> int:
-    return ContractNFTFactoryOnGnosisChain().max_supply()
diff --git a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/deploy_nft_treasury_game.py b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/deploy_nft_treasury_game.py
index 3ec790e4..1d3f3a7d 100644
--- a/prediction_market_agent/agents/microchain_agent/nft_treasury_game/deploy_nft_treasury_game.py
+++ b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/deploy_nft_treasury_game.py
@@ -25,9 +25,6 @@
     NFT_TOKEN_FACTORY,
     TREASURY_ADDRESS,
 )
-from prediction_market_agent.agents.microchain_agent.nft_treasury_game.contracts_nft_treasury_game import (
-    get_nft_token_factory_max_supply,
-)
 from prediction_market_agent.agents.microchain_agent.nft_treasury_game.messages_functions import (
     GameRoundEnd,
 )
@@ -301,7 +298,7 @@ def get_initial_system_prompt(cls) -> str:
 
 def nft_treasury_game_base_prompt(wallet_address: ChecksumAddress) -> str:
     keys = MicrochainAgentKeys()
-    n_nft_keys = get_nft_token_factory_max_supply()
+    n_nft_keys = DeployableAgentNFTGameAbstract.retrieve_total_number_of_keys()
     other_agents_keys_formatted = ", ".join(
         x.wallet_address
         for x in DEPLOYED_NFT_AGENTS