Tee pull public key support (#36)

vana-com · Oct 17, 2024 · 1f65c09 · 1f65c09
1 parent e57e769
commit 1f65c09
Show file tree

Hide file tree

Showing 8 changed files with 239 additions and 56 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vana"
-version = "0.19.0"
+version = "0.20.0"
 description = ""
 authors = ["Tim Nunamaker <[email protected]>", "Volodymyr Isai <[email protected]>", "Kahtaf Alam <[email protected]>"]
 readme = "README.md"

diff --git a/vana/__init__.py b/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.19.0"
+__version__ = "0.20.0"
 
 import rich
 

diff --git a/vana/client.py b/vana/client.py
@@ -88,13 +88,14 @@ def get_tee(self, address: str):
             return None
         return tee
 
-    def register_tee(self, url: str):
+    def register_tee(self, url: str, public_key: str):
         """
         Register a TEE compute node with the TEE Pool contract.
-        :param url:
+        :param url: URL where the TEE is reachable
+        :param public_key: Public key of the TEE node
         :return: Transaction hex, Transaction receipt
         """
-        register_fn = self.tee_pool_contract.functions.addTee(self.wallet.hotkey.address, url)
+        register_fn = self.tee_pool_contract.functions.addTee(self.wallet.hotkey.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):

diff --git a/vana/commands/satya.py b/vana/commands/satya.py
@@ -40,7 +40,7 @@ class RegisterCommand(BaseCommand):
         wallet (str): The name of the wallet to use for registration.
 
     Example usage:
-        vanacli satya register --url=https://teenode.com --wallet.name=dlp-owner --chain.network=satori
+        vanacli satya register --url=https://teenode.com --wallet.name=dlp-owner --chain.network=moksha
     """
 
     @staticmethod
@@ -52,11 +52,10 @@ 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)
-
+            tx_hash, tx_receipt = vana_client.register_tee(cli.config.url, wallet.get_hotkey_public_key())
             if tx_receipt['status'] == 1:
                 vana.__console__.print(
-                    f"[bold green]Successfully registered validator node with URL '{cli.config.url} and address {wallet.hotkey.address}'[/bold green]")
+                    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]")
                 vana.__console__.print(f"Transaction hash: {tx_hash.hex()}")
             else:
                 vana.__console__.print(

diff --git a/vana/contracts/DataRegistry.json b/vana/contracts/DataRegistry.json
@@ -314,6 +314,47 @@
     "stateMutability": "nonpayable",
     "type": "function"
   },
+  {
+    "inputs": [
+      {
+        "internalType": "string",
+        "name": "url",
+        "type": "string"
+      },
+      {
+        "internalType": "address",
+        "name": "ownerAddress",
+        "type": "address"
+      },
+      {
+        "components": [
+          {
+            "internalType": "address",
+            "name": "account",
+            "type": "address"
+          },
+          {
+            "internalType": "string",
+            "name": "key",
+            "type": "string"
+          }
+        ],
+        "internalType": "struct IDataRegistry.Permission[]",
+        "name": "permissions",
+        "type": "tuple[]"
+      }
+    ],
+    "name": "addFileWithPermissions",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "nonpayable",
+    "type": "function"
+  },
   {
     "inputs": [
       {

diff --git a/vana/contracts/RootNetworkContract.json b/vana/contracts/RootNetworkContract.json
@@ -1009,6 +1009,30 @@
     "stateMutability": "view",
     "type": "function"
   },
+  {
+    "inputs": [
+      {
+        "internalType": "uint256",
+        "name": "dlpId",
+        "type": "uint256"
+      }
+    ],
+    "name": "estimatedDlpReward",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "historyRewardEstimation",
+        "type": "uint256"
+      },
+      {
+        "internalType": "uint256",
+        "name": "stakeRewardEstimation",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  },
   {
     "inputs": [
       {
@@ -1554,6 +1578,30 @@
     "stateMutability": "nonpayable",
     "type": "function"
   },
+  {
+    "inputs": [
+      {
+        "internalType": "address",
+        "name": "stakerAddress",
+        "type": "address"
+      },
+      {
+        "internalType": "uint256",
+        "name": "dlpId",
+        "type": "uint256"
+      }
+    ],
+    "name": "unstakebleAmount",
+    "outputs": [
+      {
+        "internalType": "uint256",
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "stateMutability": "view",
+    "type": "function"
+  },
   {
     "inputs": [
       {