Introduce Vana Client for onchain interactions (#25)

pyproject.toml

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

vana/__init__.py

@@ -107,6 +107,7 @@ def turn_console_on():
 from .message import Message, TerminalInfo
 from .node_server import NodeServer
 from .node_client import NodeClient
+from .client import Client
 
 
 # Logging helpers.
@@ -144,7 +145,7 @@ def debug(on: bool = True):
 ]
 defaults = Config.merge_all(configs)
 
-banner =r"""_/\\\________/\\\_____/\\\\\\\\\_____/\\\\\_____/\\\_____/\\\\\\\\\\\___
+banner = r"""_/\\\________/\\\_____/\\\\\\\\\_____/\\\\\_____/\\\_____/\\\\\\\\\\\___
 _\/\\\_______\/\\\___/\\\\\\\\\\\\\__\/\\\\\\___\/\\\___/\\\\\\\\\\\\\__
 _\//\\\______/\\\___/\\\/////////\\\_\/\\\/\\\__\/\\\__/\\\/////////\\\_
 __\//\\\____/\\\___\/\\\_______\/\\\_\/\\\//\\\_\/\\\_\/\\\_______\/\\\_
@@ -158,6 +159,7 @@ def debug(on: bool = True):
 import sys
 import os
 
+
 def is_cli_context():
     """
     Determine if the current context is a CLI invocation by checking how the script is being executed.
@@ -173,5 +175,6 @@ def is_cli_context():
 
     return False
 
+
 if not is_cli_context():
     print(banner)

vana/chain_manager.py

@@ -129,9 +129,9 @@ def config() -> "config":
     def add_args(cls, parser: argparse.ArgumentParser, prefix: Optional[str] = None):
         prefix_str = "" if prefix is None else f"{prefix}."
         try:
-            default_network = os.getenv("OD_CHAIN_NETWORK") or "vana"
+            default_network = os.getenv("CHAIN_NETWORK") or "vana"
             default_chain_endpoint = (
-                    os.getenv("OD_CHAIN_NETWORK_ENDPOINT")
+                    os.getenv("CHAIN_NETWORK_ENDPOINT")
                     or vana.__vana_entrypoint__
             )
             parser.add_argument(

vana/client.py

@@ -0,0 +1,61 @@
+import argparse
+import copy
+import json
+import os
+from typing import Optional
+
+import vana
+from vana.contracts import contracts
+
+
+class Client:
+    @staticmethod
+    def config() -> "config":
+        parser = argparse.ArgumentParser()
+        vana.ChainManager.add_args(parser)
+        return vana.Config(parser, args=[])
+
+    def __init__(self, config: vana.Config):
+        if config is None:
+            config = self.config()
+        self.config = copy.deepcopy(config)
+
+        self.wallet = vana.Wallet(config=self.config)
+        self.chain_manager = vana.ChainManager(config=self.config)
+        self.network = self.config.chain.network
+
+        # Load contracts
+        data_registry_contract_path = os.path.join(os.path.dirname(__file__),
+                                                   "contracts/DataRegistryImplementation.json")
+        with open(data_registry_contract_path) as f:
+            self.data_registry_contract = self.chain_manager.web3.eth.contract(
+                address=contracts[self.network]["DataRegistryImplementation"],
+                abi=json.load(f)
+            )
+        tee_pool_contract_path = os.path.join(os.path.dirname(__file__),
+                                              "contracts/TeePoolImplementation.json")
+        with open(tee_pool_contract_path) as f:
+            self.tee_pool_contract = self.chain_manager.web3.eth.contract(
+                address=contracts[self.network]["TeePoolImplementation"],
+                abi=json.load(f)
+            )
+
+    # Data Registry
+
+    def get_file(self, file_id: int):
+        get_file_fn = self.data_registry_contract.functions.files(file_id)
+        file = self.chain_manager.read_contract_fn(get_file_fn)
+        (id, ownerAddress, url, addedAtBlock) = file
+        if ownerAddress == "0x0000000000000000000000000000000000000000":
+            return None
+        return file
+
+    def add_file(self, url: str, integrity_hash: Optional[str] = None):
+        add_file_fn = self.data_registry_contract.functions.addFile(url)
+        return self.chain_manager.send_transaction(add_file_fn, self.wallet.hotkey)
+
+    # TEE Pool Contract
+
+    def register_tee(self, url: str):
+        register_fn = self.tee_pool_contract.functions.addTee(self.wallet.hotkey.address, url)
+        return self.chain_manager.send_transaction(register_fn, self.wallet.hotkey)

vana/commands/satya.py

@@ -16,10 +16,12 @@
 # DEALINGS IN THE SOFTWARE.
 
 import argparse
+import traceback
+
 from rich.prompt import Prompt
+
 import vana
 from vana.commands.base_command import BaseCommand
-import traceback
 
 
 class RegisterCommand(BaseCommand):
@@ -38,7 +40,7 @@ class RegisterCommand(BaseCommand):
         wallet (str): The name of the wallet to use for registration.
 
     Example usage:
-        vanacli register satya --url=https://teenode.com --wallet=dlp-owner
+        vanacli satya register --url=https://teenode.com --wallet.name=dlp-owner --chain.network=satori
     """
 
     @staticmethod
@@ -47,62 +49,36 @@ def run(cli: "vana.cli"):
         url = cli.config.url
         vana.__console__.print(f"Registering URL with Satya: [bold]{url}[/bold]")
         try:
-            chain_manager: "vana.ChainManager" = vana.ChainManager(config=cli.config if cli.config.chain else None)
+            vana_client = vana.Client(config=cli.config)
             wallet = vana.Wallet(config=cli.config if cli.config.wallet else None)
 
-            # Connect to the contract
-            contract_address = "0x88790ffF10E952ffc13Be22a442616eAfE081594"  # TODO. Move to config
-            contract_abi = [
-                {
-                    "inputs": [
-                        {
-                            "internalType": "address",
-                            "name": "teeAddress",
-                            "type": "address"
-                        },
-                        {
-                            "internalType": "string",
-                            "name": "url",
-                            "type": "string"
-                        }
-                    ],
-                    "name": "addTee",
-                    "outputs": [],
-                    "stateMutability": "nonpayable",
-                    "type": "function"
-                }
-            ]  # TODO: Use ABIs when available
-            contract = chain_manager.web3.eth.contract(address=contract_address, abi=contract_abi)
-
-            add_tee_function = contract.functions.addTee(wallet.hotkey.address, cli.config.url)
-            tx_hash, tx_receipt = chain_manager.send_transaction(add_tee_function, wallet.hotkey)
+            tx_hash, tx_receipt = vana_client.register_tee(cli.config.url)
 
             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]")
+                vana.__console__.print(
+                    f"[bold green]Successfully registered validator node with URL '{cli.config.url} and address {wallet.hotkey.address}'[/bold green]")
                 vana.__console__.print(f"Transaction hash: {tx_hash.hex()}")
             else:
-                vana.__console__.print("[bold red]Transaction failed. Please check the contract state and try again.[/bold red]")
+                vana.__console__.print(
+                    "[bold red]Transaction failed. Please check the contract state and try again.[/bold red]")
 
         except Exception as e:
             vana.__console__.print(f"[bold red]Error:[/bold red] {str(e)}")
-            # Print stack trace
             traceback.print_exc()
-        finally:
-            if "chain_manager" in locals():
-                chain_manager.close()
-                vana.logging.debug("closing chain_manager connection")
 
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
         satya_parser = parser.add_parser(
             "register", help="Register a URL with the Satya protocol."
         )
         satya_parser.add_argument("--url", type=str, required=False, help="The URL to register.")
-        satya_parser.add_argument("--wallet.name", type=str, required=False, 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("--wallet.name", type=str, required=False,
+                                  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.")
 
     @staticmethod
     def check_config(config: "vana.Config"):
         if not config.get("url") and not config.no_prompt:
             url = Prompt.ask("Enter the URL to register")
-            config.satya.url = url
+            config.url = url