NFT functions for general agent (#582)

prediction_market_agent/agents/microchain_agent/messages_functions.py

@@ -30,7 +30,7 @@ def __call__(self, message: str) -> str:
 class SendPaidMessageToAnotherAgent(Function):
     @property
     def description(self) -> str:
-        return f"""Use {SendPaidMessageToAnotherAgent.__name__} to send a message to an another agent, given his wallet address. 
+        return f"""Use {SendPaidMessageToAnotherAgent.__name__} to send a message to an another agent, given his wallet address.
 Fee for sending the message is {TRANSACTION_MESSAGE_FEE} xDai."""
 
     @property

prediction_market_agent/agents/microchain_agent/microchain_agent.py

@@ -41,6 +41,7 @@
 from prediction_market_agent.agents.microchain_agent.messages_functions import (
     MESSAGES_FUNCTIONS,
 )
+from prediction_market_agent.agents.microchain_agent.nft_functions import NFT_FUNCTIONS
 from prediction_market_agent.agents.microchain_agent.omen_functions import (
     OMEN_FUNCTIONS,
 )
@@ -156,6 +157,9 @@ def build_agent_functions(
     if functions_config.include_messages_functions:
         functions.extend(f() for f in MESSAGES_FUNCTIONS)
 
+    if functions_config.include_nft_functions:
+        functions.extend(f() for f in NFT_FUNCTIONS)
+
     if long_term_memory:
         functions.append(
             RememberPastActions(long_term_memory=long_term_memory, model=model)

prediction_market_agent/agents/microchain_agent/microchain_agent_keys.py

@@ -10,6 +10,8 @@ class MicrochainAgentKeys(APIKeys):
     ENABLE_SOCIAL_MEDIA: bool = False
     # Double check to not spend big money during testing.
     SENDING_XDAI_CAP: float | None = OMEN_TINY_BET_AMOUNT
+    # Double check to not transfer NFTs during testing.
+    ENABLE_NFT_TRANSFER: bool = False
 
     def cap_sending_xdai(self, amount: xDai) -> xDai:
         if self.SENDING_XDAI_CAP is None:

prediction_market_agent/agents/microchain_agent/nft_functions.py

@@ -0,0 +1,97 @@
+from microchain import Function
+from prediction_market_agent_tooling.loggers import logger
+from prediction_market_agent_tooling.tools.contract import (
+    ContractOwnableERC721OnGnosisChain,
+)
+from web3 import Web3
+
+from prediction_market_agent.agents.microchain_agent.microchain_agent_keys import (
+    MicrochainAgentKeys,
+)
+
+
+class BalanceOfNFT(Function):
+    @property
+    def description(self) -> str:
+        return "Returns the number of given NFT owned by the specified address."
+
+    @property
+    def example_args(self) -> list[str]:
+        return [
+            "0xNFTAddress",
+            "0xOwnerddress",
+        ]
+
+    def __call__(
+        self,
+        nft_address: str,
+        owner_address: str,
+    ) -> int:
+        contract = ContractOwnableERC721OnGnosisChain(
+            address=Web3.to_checksum_address(nft_address)
+        )
+        balance: int = contract.balanceOf(Web3.to_checksum_address(owner_address))
+        return balance
+
+
+class OwnerOfNFT(Function):
+    @property
+    def description(self) -> str:
+        return "Returns the owner address of the specified NFT token ID."
+
+    @property
+    def example_args(self) -> list[str]:
+        return ["0xNFTAddress", "1"]
+
+    def __call__(
+        self,
+        nft_address: str,
+        token_id: int,
+    ) -> str:
+        contract = ContractOwnableERC721OnGnosisChain(
+            address=Web3.to_checksum_address(nft_address)
+        )
+        owner_address = contract.ownerOf(token_id)
+        return owner_address
+
+
+class SafeTransferFromNFT(Function):
+    @property
+    def description(self) -> str:
+        return "Transfers the specified NFT token ID from one address to another."
+
+    @property
+    def example_args(self) -> list[str]:
+        return [
+            "0xNFTAddress",
+            "0xRecipientAddress",
+            "1",
+        ]
+
+    def __call__(
+        self,
+        nft_address: str,
+        to_address: str,
+        token_id: int,
+    ) -> str:
+        keys = MicrochainAgentKeys()
+        contract = ContractOwnableERC721OnGnosisChain(
+            address=Web3.to_checksum_address(nft_address)
+        )
+        if keys.ENABLE_NFT_TRANSFER:
+            contract.safeTransferFrom(
+                api_keys=keys,
+                from_address=keys.bet_from_address,
+                to_address=Web3.to_checksum_address(to_address),
+                tokenId=token_id,
+            )
+        else:
+            logger.warning("NFT transfer is disabled in the environment.")
+        return "Token transferred successfully."
+
+
+NFT_FUNCTIONS: list[type[Function]] = [
+    BalanceOfNFT,
+    OwnerOfNFT,
+    SafeTransferFromNFT,
+]

prediction_market_agent/agents/microchain_agent/prompts.py

@@ -136,6 +136,7 @@ class FunctionsConfig(BaseModel):
     include_sending_functions: bool
     include_twitter_functions: bool
     include_messages_functions: bool
+    include_nft_functions: bool
 
     @staticmethod
     def from_system_prompt_choice(
@@ -149,6 +150,7 @@ def from_system_prompt_choice(
         include_sending_functions = False
         include_twitter_functions = False
         include_messages_functions = False
+        include_nft_functions = False
 
         if system_prompt_choice == SystemPromptChoice.JUST_BORN:
             include_learning_functions = True
@@ -171,6 +173,7 @@ def from_system_prompt_choice(
 
         elif system_prompt_choice == SystemPromptChoice.DARE_YOU_GET_MY_RESOURCES_AGENT:
             include_messages_functions = True
+            include_nft_functions = True
 
         return FunctionsConfig(
             include_trading_functions=include_trading_functions,
@@ -181,6 +184,7 @@ def from_system_prompt_choice(
             include_sending_functions=include_sending_functions,
             include_twitter_functions=include_twitter_functions,
             include_messages_functions=include_messages_functions,
+            include_nft_functions=include_nft_functions,
         )