Update register tee command, fix for network config (#51)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vana"
-version = "0.33.0"
+version = "0.34.0"
 description = ""
 authors = ["Tim Nunamaker <[email protected]>", "Volodymyr Isai <[email protected]>", "Kahtaf Alam <[email protected]>"]
 readme = "README.md"

vana/__init__.py

@@ -15,7 +15,7 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-__version__ = "0.33.0"
+__version__ = "0.34.0"
 
 import rich
 

vana/client.py

@@ -54,17 +54,26 @@ def __init__(self, config: vana.Config):
             "contracts/DataRegistry.json"
         )
         with open(data_registry_contract_path) as f:
+            data_registry_address = contracts[self.network]["DataRegistry"]
+            if hasattr(self.config, 'client') and self.config.client is not None:
+                data_registry_address = self.config.client.data_registry_contract_address or data_registry_address
+
             self.data_registry_contract = self.chain_manager.web3.eth.contract(
-                address=config.client.data_registry_contract_address or contracts[self.network]["DataRegistry"],
+                address=data_registry_address,
                 abi=json.load(f)
             )
+
         tee_pool_contract_path = os.path.join(
             os.path.dirname(__file__),
             "contracts/TeePool.json"
         )
         with open(tee_pool_contract_path) as f:
+            tee_pool_address = contracts[self.network]["TeePool"]
+            if hasattr(self.config, 'client') and self.config.client is not None:
+                tee_pool_address = self.config.client.tee_pool_contract_address or tee_pool_address
+
             self.tee_pool_contract = self.chain_manager.web3.eth.contract(
-                address=config.client.tee_pool_contract_address or contracts[self.network]["TeePool"],
+                address=tee_pool_address,
                 abi=json.load(f)
             )
 
@@ -124,14 +133,15 @@ def get_tee(self, address: str):
             return None
         return tee
 
-    def register_tee(self, url: str, public_key: str):
+    def register_tee(self, url: str, public_key: str, tee_address: str):
         """
         Register a TEE compute node with the TEE Pool contract.
-        :param url: URL where the TEE is reachable
-        :param public_key: Public key of the TEE node
-        :return: Transaction hex, Transaction receipt
+        @param url: URL where the TEE is reachable
+        @param public_key: Public key of the TEE node
+        @param tee_address: Address of the TEE to register. If not provided, uses the wallet's hotkey address.
+        @return: Transaction hex, Transaction receipt
         """
-        register_fn = self.tee_pool_contract.functions.addTee(self.wallet.hotkey.address, url, public_key)
+        register_fn = self.tee_pool_contract.functions.addTee(tee_address, url, public_key)
         return self.chain_manager.send_transaction(register_fn, self.wallet.hotkey)
 
     def add_proof(self, proof_data: ProofData, file_id: int | None = None, job_id: int | None = None):

vana/commands/satya.py

@@ -28,19 +28,19 @@ class RegisterCommand(BaseCommand):
     """
     Executes the `register` command to register a validator node on the Vana network using TEE Pool Contract.
 
-    This command allows users to register a validator node with a specified name, URL, and wallet.
+    This command allows users to register a validator node with a specified URL and wallet.
     It interacts with the TeePoolImplementation smart contract to add the validator as a TEE.
 
     Usage:
-        The command requires specifying the validator name, URL, and wallet name.
+        The command requires specifying the validator URL and wallet name.
 
     Args:
-        name (str): The name of the validator to register.
-        url (str): The URL of the validator node.
-        wallet (str): The name of the wallet to use for registration.
+        url (str): The URL of the validator node
+        tee_address (str): The hotkey address of the TEE node to register
+        wallet (str): The name of the wallet to use for registration, should have neccessary permissions
 
     Example usage:
-        vanacli satya register --url=https://teenode.com --wallet.name=dlp-owner --chain.network=moksha
+        vanacli satya register --url=https://teenode.com --wallet.name=dlp-owner --chain.network=moksha --tee_address=0x123...
     """
 
     @staticmethod
@@ -52,10 +52,21 @@ def run(cli: "vana.cli"):
             vana_client = vana.Client(config=cli.config)
             wallet = vana.Wallet(config=cli.config if cli.config.wallet else None)
 
-            tx_hash, tx_receipt = vana_client.register_tee(cli.config.url, wallet.get_hotkey_public_key())
+            tee_address = cli.config.get('tee_address', wallet.hotkey.address)
+
+            tx_hash, tx_receipt = vana_client.register_tee(
+                url=cli.config.url,
+                public_key=wallet.get_hotkey_public_key(),
+                tee_address=tee_address
+            )
+
             if tx_receipt['status'] == 1:
                 vana.__console__.print(
-                    f"[bold green]Successfully registered validator node with URL '{cli.config.url}', address '{wallet.hotkey.address}' and public key '{wallet.get_hotkey_public_key()}'[/bold green]")
+                    f"[bold green]Successfully registered validator node with:"
+                    f"\n- URL: '{cli.config.url}'"
+                    f"\n- TEE Address: '{tee_address}'"
+                    f"\n- Public Key: '{wallet.get_hotkey_public_key()}'[/bold green]"
+                )
                 vana.__console__.print(f"Transaction hash: {tx_hash.hex()}")
             else:
                 vana.__console__.print(
@@ -75,6 +86,8 @@ def add_args(parser: argparse.ArgumentParser):
                                   help="The name of the wallet to use for registration.")
         satya_parser.add_argument("--chain.network", type=str, required=False,
                                   help="The network to use for registration.")
+        satya_parser.add_argument("--tee_address", type=str, required=False,
+                                  help="The hotkey address of the TEE node to register.")
 
     @staticmethod
     def check_config(config: "vana.Config"):