feat: token balances

src/actions/balance.ts

@@ -1,6 +1,6 @@
 import { PublicKey } from "@solana/web3.js";
-import { Action } from "../types/action";
-import { SolanaAgentKit } from "../agent";
+import type { Action } from "../types/action";
+import type { SolanaAgentKit } from "../agent";
 import { z } from "zod";
 import { get_balance } from "../tools";
 
@@ -54,6 +54,7 @@ const balanceAction: Action = {
     return {
       status: "success",
       balance: balance,
+      token: input.tokenAddress || "SOL",
     };
   },
 };

src/actions/index.ts

@@ -30,9 +30,11 @@ import launchPumpfunTokenAction from "./launchPumpfunToken";
 import getWalletAddressAction from "./getWalletAddress";
 import flashOpenTradeAction from "./flashOpenTrade";
 import flashCloseTradeAction from "./flashCloseTrade";
+import tokenBalancesAction from "./tokenBalances";
 
 export const ACTIONS = {
   WALLET_ADDRESS_ACTION: getWalletAddressAction,
+  TOKEN_BALANCES_ACTION: tokenBalancesAction,
   DEPLOY_TOKEN_ACTION: deployTokenAction,
   BALANCE_ACTION: balanceAction,
   TRANSFER_ACTION: transferAction,

src/actions/tokenBalances.ts

@@ -0,0 +1,80 @@
+import { PublicKey } from "@solana/web3.js";
+import type { Action } from "../types/action";
+import type { SolanaAgentKit } from "../agent";
+import { z } from "zod";
+import { get_token_balance } from "../tools/get_token_balances";
+
+const tokenBalancesAction: Action = {
+  name: "TOKEN_BALANCE_ACTION",
+  similes: [
+    "check token balances",
+    "get wallet token balances",
+    "view token balances",
+    "show token balances",
+    "check token balance",
+  ],
+  description: `Get the token balances of a Solana wallet.
+  If you want to get the balance of your wallet, you don't need to provide the wallet address.`,
+  examples: [
+    [
+      {
+        input: {},
+        output: {
+          status: "success",
+          balance: {
+            sol: 100,
+            tokens: [
+              {
+                tokenAddress: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
+                name: "USD Coin",
+                symbol: "USDC",
+                balance: 100,
+                decimals: 9,
+              },
+            ],
+          },
+        },
+        explanation: "Get token balances of the wallet",
+      },
+    ],
+    [
+      {
+        input: {
+          walletAddress: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
+        },
+        output: {
+          status: "success",
+          balance: {
+            sol: 100,
+            tokens: [
+              {
+                tokenAddress: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
+                name: "USD Coin",
+                symbol: "USDC",
+                balance: 100,
+                decimals: 9,
+              },
+            ],
+          },
+        },
+        explanation: "Get address token balance",
+      },
+    ],
+  ],
+  schema: z.object({
+    walletAddress: z.string().optional(),
+  }),
+  handler: async (agent: SolanaAgentKit, input: Record<string, any>) => {
+    const balance = await get_token_balance(
+      agent,
+      input.tokenAddress && new PublicKey(input.tokenAddress),
+    );
+
+    return {
+      status: "success",
+      balance: balance,
+    };
+  },
+};
+
+export default tokenBalancesAction;

src/tools/get_balance.ts

@@ -1,7 +1,5 @@
-import { LAMPORTS_PER_SOL, PublicKey } from "@solana/web3.js";
-import { SolanaAgentKit } from "../index";
-import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
-import { getTokenMetadata } from "../utils/tokenMetadata";
+import { LAMPORTS_PER_SOL, type PublicKey } from "@solana/web3.js";
+import type { SolanaAgentKit } from "../index";
 
 /**
  * Get the balance of SOL or an SPL token for the agent's wallet
@@ -12,51 +10,12 @@ import { getTokenMetadata } from "../utils/tokenMetadata";
 export async function get_balance(
   agent: SolanaAgentKit,
   token_address?: PublicKey,
-): Promise<
-  | number
-  | {
-      sol: number;
-      tokens: Array<{
-        tokenAddress: string;
-        name: string;
-        symbol: string;
-        balance: number;
-        decimals: number;
-      }>;
-    }
-> {
+): Promise<number> {
   if (!token_address) {
-    const [lamportsBalance, tokenAccountData] = await Promise.all([
-      agent.connection.getBalance(agent.wallet_address),
-      agent.connection.getParsedTokenAccountsByOwner(agent.wallet_address, {
-        programId: TOKEN_PROGRAM_ID,
-      }),
-    ]);
-
-    const removedZeroBalance = tokenAccountData.value.filter(
-      (v) => v.account.data.parsed.info.tokenAmount.uiAmount !== 0,
+    return (
+      (await agent.connection.getBalance(agent.wallet_address)) /
+      LAMPORTS_PER_SOL
     );
-
-    const tokenBalances = await Promise.all(
-      removedZeroBalance.map(async (v) => {
-        const mint = v.account.data.parsed.info.mint;
-        const mintInfo = await getTokenMetadata(agent.connection, mint);
-        return {
-          tokenAddress: mint,
-          name: mintInfo.name ?? "",
-          symbol: mintInfo.symbol ?? "",
-          balance: v.account.data.parsed.info.tokenAmount.uiAmount as number,
-          decimals: v.account.data.parsed.info.tokenAmount.decimals as number,
-        };
-      }),
-    );
-
-    const solBalance = lamportsBalance / LAMPORTS_PER_SOL;
-
-    return {
-      sol: solBalance,
-      tokens: tokenBalances,
-    };
   }
 
   const token_account =

src/tools/get_token_balances.ts

@@ -0,0 +1,65 @@
+import { LAMPORTS_PER_SOL, type PublicKey } from "@solana/web3.js";
+import type { SolanaAgentKit } from "../index";
+import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import { getTokenMetadata } from "../utils/tokenMetadata";
+
+/**
+ * Get the token balances of a Solana wallet
+ * @param agent - SolanaAgentKit instance
+ * @param token_address - Optional SPL token mint address. If not provided, returns SOL balance
+ * @returns Promise resolving to the balance as an object containing sol balance and token balances with their respective mints, symbols, names and decimals
+ */
+export async function get_token_balance(
+  agent: SolanaAgentKit,
+  token_address?: PublicKey,
+): Promise<
+  | number
+  | {
+      sol: number;
+      tokens: Array<{
+        tokenAddress: string;
+        name: string;
+        symbol: string;
+        balance: number;
+        decimals: number;
+      }>;
+    }
+> {
+  if (!token_address) {
+    const [lamportsBalance, tokenAccountData] = await Promise.all([
+      agent.connection.getBalance(agent.wallet_address),
+      agent.connection.getParsedTokenAccountsByOwner(agent.wallet_address, {
+        programId: TOKEN_PROGRAM_ID,
+      }),
+    ]);
+
+    const removedZeroBalance = tokenAccountData.value.filter(
+      (v) => v.account.data.parsed.info.tokenAmount.uiAmount !== 0,
+    );
+
+    const tokenBalances = await Promise.all(
+      removedZeroBalance.map(async (v) => {
+        const mint = v.account.data.parsed.info.mint;
+        const mintInfo = await getTokenMetadata(agent.connection, mint);
+        return {
+          tokenAddress: mint,
+          name: mintInfo.name ?? "",
+          symbol: mintInfo.symbol ?? "",
+          balance: v.account.data.parsed.info.tokenAmount.uiAmount as number,
+          decimals: v.account.data.parsed.info.tokenAmount.decimals as number,
+        };
+      }),
+    );
+
+    const solBalance = lamportsBalance / LAMPORTS_PER_SOL;
+
+    return {
+      sol: solBalance,
+      tokens: tokenBalances,
+    };
+  }
+
+  const token_account =
+    await agent.connection.getTokenAccountBalance(token_address);
+  return token_account.value.uiAmount || 0;
+}

test/agent_sdks/vercel_ai.ts

@@ -95,7 +95,6 @@ async function runChatMode() {
   );
 
   const tools = createVercelAITools(solanaAgent);
-  console.log(tools);
 
   const rl = readline.createInterface({
     input: process.stdin,