Updates swap stake

Author: ibraheem-abe

bittensor/core/subtensor.py

@@ -2769,6 +2769,9 @@ def swap_stake(
         amount: Balance,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
+        safe_staking: bool = False,
+        allow_partial_stake: bool = False,
+        rate_threshold: float = 0.005,
     ) -> bool:
         """
         Moves stake between subnets while keeping the same coldkey-hotkey pair ownership.
@@ -2782,9 +2785,26 @@ def swap_stake(
             amount (Union[Balance, float]): The amount to swap.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
+            safe_staking (bool): If true, enables price safety checks to protect against price impact. The swap
+                will only execute if the price ratio between subnets doesn't exceed the rate threshold.
+                Default is False.
+            allow_partial_stake (bool): If true and safe_staking is enabled, allows partial stake swaps when
+                the full amount would exceed the price threshold. If false, the entire swap fails if it would
+                exceed the threshold. Default is False.
+            rate_threshold (float): The maximum allowed increase in the price ratio between subnets
+                (origin_price/destination_price). For example, 0.005 = 0.5% maximum increase. Only used
+                when safe_staking is True. Default is 0.005.
+
 
         Returns:
             success (bool): True if the extrinsic was successful.
+
+        The price ratio for swap_stake in safe mode is calculated as: origin_subnet_price / destination_subnet_price
+        When safe_staking is enabled, the swap will only execute if:
+            - With allow_partial_stake=False: The entire swap amount can be executed without the price ratio
+            increasing more than rate_threshold
+            - With allow_partial_stake=True: A partial amount will be swapped up to the point where the
+            price ratio would increase by rate_threshold
         """
         amount = check_and_convert_to_balance(amount)
         return swap_stake_extrinsic(
@@ -2796,6 +2816,9 @@ def swap_stake(
             amount=amount,
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
+            safe_staking=safe_staking,
+            allow_partial_stake=allow_partial_stake,
+            rate_threshold=rate_threshold,
         )
 
     def transfer(