Merge pull request #632 from opentensor/feat/auto-staking

Feat/auto staking

Author: ibraheem-abe

bittensor_cli/cli.py

@@ -72,6 +72,7 @@
     prompt_position_id,
 )
 from bittensor_cli.src.commands.stake import (
+    auto_staking as auto_stake,
     children_hotkeys,
     list as list_stake,
     move as move_stake,
@@ -935,6 +936,12 @@ def __init__(self):
         self.stake_app.command(
             "add", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"]
         )(self.stake_add)
+        self.stake_app.command(
+            "auto", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"]
+        )(self.get_auto_stake)
+        self.stake_app.command(
+            "set-auto", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"]
+        )(self.set_auto_stake)
         self.stake_app.command(
             "remove", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"]
         )(self.stake_remove)
@@ -3596,6 +3603,137 @@ def wallet_swap_coldkey(
             )
         )
 
+    def get_auto_stake(
+        self,
+        network: Optional[list[str]] = Options.network,
+        wallet_name: Optional[str] = Options.wallet_name,
+        wallet_path: Optional[str] = Options.wallet_path,
+        coldkey_ss58=typer.Option(
+            None,
+            "--ss58",
+            "--coldkey_ss58",
+            "--coldkey.ss58_address",
+            "--coldkey.ss58",
+            help="Coldkey address of the wallet",
+        ),
+        quiet: bool = Options.quiet,
+        verbose: bool = Options.verbose,
+        json_output: bool = Options.json_output,
+    ):
+        """Display auto-stake destinations for a wallet across all subnets."""
+
+        self.verbosity_handler(quiet, verbose, json_output)
+
+        wallet = None
+        if coldkey_ss58:
+            if not is_valid_ss58_address(coldkey_ss58):
+                print_error("You entered an invalid ss58 address")
+                raise typer.Exit()
+        else:
+            if wallet_name:
+                coldkey_or_ss58 = wallet_name
+            else:
+                coldkey_or_ss58 = Prompt.ask(
+                    "Enter the [blue]wallet name[/blue] or [blue]coldkey ss58 address[/blue]",
+                    default=self.config.get("wallet_name") or defaults.wallet.name,
+                )
+            if is_valid_ss58_address(coldkey_or_ss58):
+                coldkey_ss58 = coldkey_or_ss58
+            else:
+                wallet_name = coldkey_or_ss58 if coldkey_or_ss58 else wallet_name
+                wallet = self.wallet_ask(
+                    wallet_name,
+                    wallet_path,
+                    None,
+                    ask_for=[WO.NAME, WO.PATH],
+                    validate=WV.WALLET,
+                )
+
+        return self._run_command(
+            auto_stake.show_auto_stake_destinations(
+                wallet,
+                self.initialize_chain(network),
+                coldkey_ss58=coldkey_ss58,
+                json_output=json_output,
+                verbose=verbose,
+            )
+        )
+
+    def set_auto_stake(
+        self,
+        network: Optional[list[str]] = Options.network,
+        wallet_name: Optional[str] = Options.wallet_name,
+        wallet_path: Optional[str] = Options.wallet_path,
+        netuid: Optional[int] = Options.netuid_not_req,
+        quiet: bool = Options.quiet,
+        verbose: bool = Options.verbose,
+        prompt: bool = Options.prompt,
+        wait_for_inclusion: bool = Options.wait_for_inclusion,
+        wait_for_finalization: bool = Options.wait_for_finalization,
+        json_output: bool = Options.json_output,
+    ):
+        """Set the auto-stake destination hotkey for a coldkey."""
+
+        self.verbosity_handler(quiet, verbose, json_output)
+
+        wallet = self.wallet_ask(
+            wallet_name,
+            wallet_path,
+            None,
+            ask_for=[WO.NAME, WO.PATH],
+            validate=WV.WALLET,
+        )
+
+        if netuid is None:
+            netuid = IntPrompt.ask(
+                "Enter the [blue]netuid[/blue] to configure",
+                default=defaults.netuid,
+            )
+        validate_netuid(netuid)
+
+        hotkey_prompt = Prompt.ask(
+            "Enter the [blue]hotkey ss58 address[/blue] to auto-stake to "
+            "[dim](Press Enter to view delegates)[/dim]",
+            default="",
+            show_default=False,
+        ).strip()
+
+        if not hotkey_prompt:
+            selected_hotkey = self._run_command(
+                subnets.show(
+                    subtensor=self.initialize_chain(network),
+                    netuid=netuid,
+                    sort=False,
+                    max_rows=20,
+                    prompt=False,
+                    delegate_selection=True,
+                ),
+                exit_early=False,
+            )
+            if not selected_hotkey:
+                print_error("No delegate selected. Exiting.")
+                return
+            hotkey_ss58 = selected_hotkey
+        else:
+            hotkey_ss58 = hotkey_prompt
+
+        if not is_valid_ss58_address(hotkey_ss58):
+            print_error("You entered an invalid hotkey ss58 address")
+            return
+
+        return self._run_command(
+            auto_stake.set_auto_stake_destination(
+                wallet,
+                self.initialize_chain(network),
+                netuid,
+                hotkey_ss58,
+                wait_for_inclusion=wait_for_inclusion,
+                wait_for_finalization=wait_for_finalization,
+                prompt_user=prompt,
+                json_output=json_output,
+            )
+        )
+
     def stake_list(
         self,
         network: Optional[list[str]] = Options.network,

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -225,6 +225,30 @@ async def get_stake_for_coldkey(
         stakes: list[StakeInfo] = StakeInfo.list_from_any(result)
         return [stake for stake in stakes if stake.stake > 0]
 
+    async def get_auto_stake_destinations(
+        self,
+        coldkey_ss58: str,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> dict[int, str]:
+        """Retrieve auto-stake destinations configured for a coldkey."""
+
+        query = await self.substrate.query_map(
+            module="SubtensorModule",
+            storage_function="AutoStakeDestination",
+            params=[coldkey_ss58],
+            block_hash=block_hash,
+            reuse_block_hash=reuse_block,
+        )
+
+        destinations: dict[int, str] = {}
+        async for netuid, destination in query:
+            hotkey_ss58 = decode_account_id(destination.value[0])
+            if hotkey_ss58:
+                destinations[int(netuid)] = hotkey_ss58
+
+        return destinations
+
     async def get_stake_for_coldkey_and_hotkey(
         self,
         hotkey_ss58: str,