chore: remove get all tokens balances (#3)

OriginTrail · Jan 20, 2025 · 8aebe1c · 8aebe1c
1 parent 7811488
commit 8aebe1c
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 75 deletions.
diff --git a/packages/plugin-bnb/src/actions/getBalance.ts b/packages/plugin-bnb/src/actions/getBalance.ts
@@ -8,7 +8,7 @@ import {
     type Memory,
     type State,
 } from "@elizaos/core";
-import { getToken, getTokens, getTokenBalances, ChainId } from "@lifi/sdk";
+import { getToken } from "@lifi/sdk";
 
 import {
     bnbWalletProvider,
@@ -17,7 +17,6 @@ import {
 } from "../providers/wallet";
 import { getBalanceTemplate } from "../templates";
 import type {
-    Balance,
     GetBalanceParams,
     GetBalanceResponse,
     SupportedChain,
@@ -46,47 +45,37 @@ export class GetBalanceAction {
             let resp: GetBalanceResponse = {
                 chain,
                 address: address!,
-                balances: [],
             };
 
-            // If no specific token is requested, get all token balances
-            if (!token) {
-                this.walletProvider.configureLiFiSdk(chain);
-                const balances = await this.getTokenBalances(chainId, address!);
-                resp.balances = balances;
-            } else {
-                // If specific token is requested and it's not the native token
-                if (token.toLowerCase() !== nativeSymbol.toLowerCase()) {
-                    let balance: string;
-                    if (token.startsWith("0x")) {
-                        balance = await this.getERC20TokenBalance(
-                            chain,
-                            address!,
-                            token as `0x${string}`
-                        );
-                    } else {
-                        this.walletProvider.configureLiFiSdk(chain);
-                        const tokenInfo = await getToken(chainId, token);
-                        balance = await this.getERC20TokenBalance(
-                            chain,
-                            address!,
-                            tokenInfo.address as `0x${string}`
-                        );
-                    }
-
-                    resp.balances = [{ token, balance }];
+            // If ERC20 token is requested
+            if (token.toLowerCase() !== nativeSymbol.toLowerCase()) {
+                let amount: string;
+                if (token.startsWith("0x")) {
+                    amount = await this.getERC20TokenBalance(
+                        chain,
+                        address!,
+                        token as `0x${string}`
+                    );
                 } else {
-                    // If native token is requested
-                    const nativeBalanceWei = await this.walletProvider
-                        .getPublicClient(chain)
-                        .getBalance({ address: address! });
-                    resp.balances = [
-                        {
-                            token: nativeSymbol,
-                            balance: formatEther(nativeBalanceWei),
-                        },
-                    ];
+                    this.walletProvider.configureLiFiSdk(chain);
+                    const tokenInfo = await getToken(chainId, token);
+                    amount = await this.getERC20TokenBalance(
+                        chain,
+                        address!,
+                        tokenInfo.address as `0x${string}`
+                    );
                 }
+
+                resp.balance = { token, amount };
+            } else {
+                // If native token is requested
+                const nativeBalanceWei = await this.walletProvider
+                    .getPublicClient(chain)
+                    .getBalance({ address: address! });
+                resp.balance = {
+                    token: nativeSymbol,
+                    amount: formatEther(nativeBalanceWei),
+                };
             }
 
             return resp;
@@ -118,22 +107,6 @@ export class GetBalanceAction {
         return formatUnits(balance, decimals);
     }
 
-    async getTokenBalances(
-        chainId: ChainId,
-        address: Address
-    ): Promise<Balance[]> {
-        const tokensResponse = await getTokens();
-        const tokens = tokensResponse.tokens[chainId];
-
-        const tokenBalances = await getTokenBalances(address, tokens);
-        return tokenBalances
-            .filter((balance) => balance.amount && balance.amount !== 0n)
-            .map((balance) => ({
-                token: balance.symbol,
-                balance: formatUnits(balance.amount!, balance.decimals),
-            }));
-    }
-
     async validateAndNormalizeParams(params: GetBalanceParams): Promise<void> {
         if (!params.address) {
             params.address = this.walletProvider.getAddress();
@@ -143,11 +116,15 @@ export class GetBalanceAction {
             );
         }
 
+        if (!params.token || params.token === "") {
+            params.token = "BNB";
+        }
+
         if (params.chain != "bsc") {
             // if token contract address is not provided, only BSC mainnet is supported
-            if (!(params.token && params.token.startsWith("0x"))) {
+            if (!params.token.startsWith("0x")) {
                 throw new Error(
-                    "If token contract address is not provided, only BSC mainnet is supported"
+                    "Only BSC mainnet is supported for querying balance by token symbol"
                 );
             }
         }
@@ -196,10 +173,10 @@ export const getBalanceAction = {
             const getBalanceResp = await action.getBalance(getBalanceOptions);
             if (callback) {
                 let text = `No balance found for ${getBalanceOptions.address} on ${getBalanceOptions.chain}`;
-                if (getBalanceResp.balances.length > 0) {
-                    text = `Balance of ${getBalanceResp.address} on ${getBalanceResp.chain}:\n${getBalanceResp.balances
-                        .map(({ token, balance }) => `${token}: ${balance}`)
-                        .join("\n")}`;
+                if (getBalanceResp.balance) {
+                    text = `Balance of ${getBalanceResp.address} on ${getBalanceResp.chain}:\n${
+                        getBalanceResp.balance.token
+                    }: ${getBalanceResp.balance.amount}`;
                 }
                 callback({
                     text,

diff --git a/packages/plugin-bnb/src/templates/index.ts b/packages/plugin-bnb/src/templates/index.ts
@@ -7,7 +7,7 @@ export const getBalanceTemplate = `Given the recent messages and wallet informat
 Extract the following information about the requested check balance:
 - Chain to execute on. Must be one of ["bsc", "bscTestnet", "opBNB", "opBNBTestnet"]. Default is "bsc".
 - Address to check balance for. Optional, must be a valid Ethereum address starting with "0x" or a web3 domain name. If not provided, use the BNB chain Wallet Address.
-- Token symbol or address. Optional. The address must be a valid Ethereum address starting with "0x".
+- Token symbol or address. Could be a token symbol or address. If the address is provided, it must be a valid Ethereum address starting with "0x". Default is "BNB".
 If any field is not provided, use the default value. If no default value is specified, use null.
 
 Respond with a JSON markdown block containing only the extracted values. Use null for any values that cannot be determined:
@@ -16,7 +16,7 @@ Respond with a JSON markdown block containing only the extracted values. Use nul
 {
     "chain": SUPPORTED_CHAINS,
     "address": string | null,
-    "token": string | null
+    "token": string
 }
 \`\`\`
 `;

diff --git a/packages/plugin-bnb/src/tests/getBalance.test.ts b/packages/plugin-bnb/src/tests/getBalance.test.ts
@@ -29,7 +29,7 @@ describe("GetBalance Action", () => {
                 token: "BNB",
             };
             const resp = await ga.getBalance(input);
-            console.log("BNB balance", resp.balances[0]);
+            console.log("BNB balance", resp.balance);
         });
 
         it("get USDC balance", async () => {
@@ -39,16 +39,20 @@ describe("GetBalance Action", () => {
                 token: "USDC",
             };
             const resp = await ga.getBalance(input);
-            console.log("USDC balance", resp.balances[0]);
+            console.log("USDC balance", resp.balance);
         });
 
-        it("get all token balances", async () => {
+        it("get balance by token contract address", async () => {
             const input: GetBalanceParams = {
                 chain: "bsc",
                 address: account.address,
+                token: "0x55d398326f99059ff775485246999027b3197955",
             };
             const resp = await ga.getBalance(input);
-            console.log("token balances", resp.balances);
-        }, 50000);
+            console.log(
+                "0x55d398326f99059ff775485246999027b3197955 balance",
+                resp.balance
+            );
+        });
     });
 });
diff --git a/packages/plugin-bnb/src/types/index.ts b/packages/plugin-bnb/src/types/index.ts
@@ -3,16 +3,11 @@ import type { Address, Hash } from "viem";
 export type SupportedChain = "bsc" | "bscTestnet" | "opBNB" | "opBNBTestnet";
 export type StakeAction = "deposit" | "withdraw" | "claim";
 
-export interface Balance {
-    token: string;
-    balance: string;
-}
-
 // Action parameters
 export interface GetBalanceParams {
     chain: SupportedChain;
     address?: Address;
-    token?: string;
+    token: string;
 }
 
 export interface TransferParams {
@@ -55,7 +50,7 @@ export interface FaucetParams {
 export interface GetBalanceResponse {
     chain: SupportedChain;
     address: Address;
-    balances: Balance[];
+    balance?: { token: string; amount: string };
 }
 
 export interface TransferResponse {