Feature: InjectiveStaking add new compound_rewards function

injective_functions/staking/__init__.py

@@ -1,6 +1,10 @@
+import asyncio
 from decimal import Decimal
+
+from pyinjective.constant import ADDITIONAL_CHAIN_FORMAT_DECIMALS, INJ_DENOM
+
 from injective_functions.base import InjectiveBase
-from typing import Dict, List
+from typing import Dict
 
 
 """This class handles all account transfer within the account"""
@@ -19,3 +23,87 @@ async def stake_tokens(self, validator_address: str, amount: str) -> Dict:
             amount=float(amount),
         )
         return await self.chain_client.build_and_broadcast_tx(msg)
+
+    async def compound_rewards(self, validator_address: str) -> Dict:
+        """
+        Compounds staking rewards by withdrawing them and restaking.
+        :param validator_address: The validator's address
+        :return: Transaction result
+        """
+        try:
+            if not validator_address.startswith("injvaloper"):
+                raise ValueError("Invalid validator address format")
+
+            # Step 1: Fetch the initial INJ balance
+            balance_response = await self.chain_client.client.get_bank_balance(
+                address=self.chain_client.address.to_acc_bech32(),
+                denom=INJ_DENOM
+            )
+            initial_balance = Decimal(balance_response.balance.amount)
+
+            # Step 2: Withdraw rewards
+            withdraw_msg = self.chain_client.composer.msg_withdraw_delegator_reward(
+                delegator_address=self.chain_client.address.to_acc_bech32(),
+                validator_address=validator_address,
+            )
+            withdraw_response = await self.chain_client.build_and_broadcast_tx(withdraw_msg)
+
+            # Step 3: Wait for the balance to update
+            updated_balance = await self.wait_for_balance_update(old_balance=initial_balance, denom=INJ_DENOM)
+
+            # Step 4: Calculate the withdrawn rewards
+            rewards_to_stake = updated_balance - initial_balance
+            if rewards_to_stake < 0:
+                return {
+                    "success": False,
+                    "error": f"Rewards ({rewards_to_stake}) are lower than gas fees, resulting in a negative net reward."
+                }
+
+            if rewards_to_stake == 0:
+                return {"success": False, "error": "No rewards available to compound."}
+
+            # Step 5: Restake the rewards
+            delegate_msg = self.chain_client.composer.MsgDelegate(
+                delegator_address=self.chain_client.address.to_acc_bech32(),
+                validator_address=validator_address,
+                amount=rewards_to_stake / Decimal(f"1e{ADDITIONAL_CHAIN_FORMAT_DECIMALS}"),
+            )
+            delegate_response = await self.chain_client.build_and_broadcast_tx(delegate_msg)
+
+            return {
+                "success": True,
+                "withdraw_response": withdraw_response,
+                "delegate_response": delegate_response,
+            }
+
+        except (TimeoutError, ValueError) as e:
+            return {"success": False, "error": str(e)}
+
+    async def wait_for_balance_update(
+        self,
+        old_balance: Decimal,
+        denom: str,
+        timeout: int = 10,
+        interval: int = 1
+    ) -> Decimal:
+        """
+        Waits for the balance to update after a transaction.
+        :param old_balance: Previous balance to compare against
+        :param denom: Denomination of the token (e.g., "inj")
+        :param timeout: Total time to wait (in seconds)
+        :param interval: Time between balance checks (in seconds)
+        :return: Updated balance
+        """
+        if interval <= 0:
+            raise ValueError("Interval must be greater than zero.")
+
+        for _ in range(timeout // interval):
+            balance_response = await self.chain_client.client.get_bank_balance(
+                address=self.chain_client.address.to_acc_bech32(),
+                denom=denom
+            )
+            updated_balance = Decimal(balance_response.balance.amount)
+            if updated_balance != old_balance:
+                return updated_balance
+            await asyncio.sleep(interval)
+        raise TimeoutError("Balance did not update within the timeout period.")

injective_functions/staking/staking_schema.json

@@ -17,6 +17,20 @@
           },
           "required": ["validator_address", "amount"]
         }
+      },
+      {
+        "name": "compound_rewards",
+        "description": "Automatically reinvest your staking rewards with a specific validator to increase your staked amount.",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "validator_address": {
+              "type": "string",
+              "description": "Validator address you want to compound your rewards with."
+            }
+          },
+          "required": ["validator_address"]
+        }
       }
     ]
   }

injective_functions/utils/function_helper.py

@@ -45,6 +45,7 @@ class InjectiveFunctionMapper:
         "query_total_supply": ("bank", "query_total_supply"),
         # Staking functions
         "stake_tokens": ("staking", "stake_tokens"),
+        "compound_rewards": ("staking", "compound_rewards"),
         # Auction functions
         "send_bid_auction": ("auction", "send_bid_auction"),
         "fetch_auctions": ("auction", "fetch_auctions"),