From cc795934131653375a8bc33a56896b7e167c2084 Mon Sep 17 00:00:00 2001
From: Peter Jung <peter@jung.ninja>
Date: Tue, 17 Dec 2024 08:35:17 +0100
Subject: [PATCH] more tweaks

---
 .../app_nft_treasury_game.py                  |  5 +++-
 .../contracts_nft_treasury_game.py            | 24 +++++++++++++++++++
 .../deploy_nft_treasury_game.py               |  8 +++++--
 3 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py

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 dbb0a2c2..4ae8d208 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
@@ -21,6 +21,9 @@
     NFT_TOKEN_FACTORY,
     TREASURY_SAFE_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,
@@ -246,7 +249,7 @@ def show_treasury_part() -> None:
     treasury_xdai_balance = get_balances(TREASURY_SAFE_ADDRESS).xdai
     st.markdown(
         f"""### Treasury
-Currently holds <span style='font-size: 1.1em;'><strong>{treasury_xdai_balance:.2f} xDAI</strong></span>.""",
+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.""",
         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
new file mode 100644
index 00000000..ca3d3480
--- /dev/null
+++ b/prediction_market_agent/agents/microchain_agent/nft_treasury_game/contracts_nft_treasury_game.py
@@ -0,0 +1,24 @@
+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,
+)
+
+
+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
+
+
+@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 4a374814..85bc02cb 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
@@ -14,6 +14,9 @@
     NFT_TOKEN_FACTORY,
     TREASURY_SAFE_ADDRESS,
 )
+from prediction_market_agent.agents.microchain_agent.nft_treasury_game.contracts_nft_treasury_game import (
+    get_nft_token_factory_max_supply,
+)
 
 
 class DeployableAgentNFTGameAbstract(DeployableMicrochainAgentAbstract):
@@ -184,6 +187,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()
     other_agents_keys_formatted = ", ".join(
         x.wallet_address
         for x in DEPLOYED_NFT_AGENTS
@@ -194,8 +198,8 @@ def nft_treasury_game_base_prompt(wallet_address: ChecksumAddress) -> str:
 - You participate in the securing of the NFT key to a treasury.
 - Your wallet address is {wallet_address}.
 - Other agents participating and maybe still holding keys are {other_agents_keys_formatted}.
-- Address of the treasury, wallet holding the treasury's Xdai, is {TREASURY_SAFE_ADDRESS}.
-- Address of the NFT contract is {NFT_TOKEN_FACTORY}, there are 5 keys, with token_id 0, 1, 2, 3, 4. 
+- Address of the treasury, wallet holding the treasury's xDai, is {TREASURY_SAFE_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.
 - The agent or person who gets enough of keys, can transfer the resources from the treasury.