get_solana_token_balance_by_mint_address

goat-sdk · Dec 9, 2024 · e980ef8 · e980ef8
1 parent 9288b86
commit e980ef8
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/typescript/packages/plugins/spl-token/src/methods/balance.ts b/typescript/packages/plugins/spl-token/src/methods/balance.ts
@@ -1,9 +1,8 @@
 import { getAssociatedTokenAddressSync } from "@solana/spl-token";
 import { type Connection, PublicKey } from "@solana/web3.js";
-import type { NetworkSpecificToken } from "../tokens";
 
-export async function balanceOf(connection: Connection, walletAddress: string, token: NetworkSpecificToken) {
-    const tokenAccount = getAssociatedTokenAddressSync(new PublicKey(token.mintAddress), new PublicKey(walletAddress));
+export async function balanceOf(connection: Connection, walletAddress: string, tokenAddress: string) {
+    const tokenAccount = getAssociatedTokenAddressSync(new PublicKey(tokenAddress), new PublicKey(walletAddress));
     const balance = await connection.getTokenAccountBalance(tokenAccount);
     return balance;
 }
diff --git a/typescript/packages/plugins/spl-token/src/parameters.ts b/typescript/packages/plugins/spl-token/src/parameters.ts
@@ -8,3 +8,7 @@ export const transferParametersSchema = z.object({
     to: z.string().describe("The address to transfer the token to"),
     amount: z.string().describe("The amount of tokens to transfer"),
 });
+
+export const getTokenBalanceByMintAddressParametersSchema = getBalanceParametersSchema.extend({
+    mintAddress: z.string().describe("The mint address of the token to get the balance of"),
+});
diff --git a/typescript/packages/plugins/spl-token/src/utils/getTools.ts b/typescript/packages/plugins/spl-token/src/utils/getTools.ts
@@ -3,7 +3,11 @@ import type { Connection } from "@solana/web3.js";
 import type { z } from "zod";
 import { balanceOf } from "../methods/balance";
 import { transfer } from "../methods/transfer";
-import { getBalanceParametersSchema, transferParametersSchema } from "../parameters";
+import {
+    getBalanceParametersSchema,
+    getTokenBalanceByMintAddressParametersSchema,
+    transferParametersSchema,
+} from "../parameters";
 import type { NetworkSpecificToken } from "../tokens";
 
 export function getTools(
@@ -14,15 +18,15 @@ export function getTools(
 
     for (const token of tokenList) {
         const balanceTool: DeferredTool<SolanaWalletClient> = {
-            name: `get_${token.symbol}_balance`,
+            name: `get_solana_${token.symbol}_balance`,
             description: `This {{tool}} gets the balance of ${token.symbol}`,
             parameters: getBalanceParametersSchema,
             method: async (walletClient: SolanaWalletClient, parameters: z.infer<typeof getBalanceParametersSchema>) =>
-                balanceOf(connection, parameters.wallet, token),
+                balanceOf(connection, parameters.wallet, token.mintAddress),
         };
 
         const transferTool: DeferredTool<SolanaWalletClient> = {
-            name: `transfer_${token.symbol}`,
+            name: `transfer_solana_${token.symbol}`,
             description: `This {{tool}} transfers ${token.symbol}`,
             parameters: transferParametersSchema,
             method: async (walletClient: SolanaWalletClient, parameters: z.infer<typeof transferParametersSchema>) =>
@@ -32,5 +36,15 @@ export function getTools(
         tools.push(balanceTool, transferTool);
     }
 
+    tools.push({
+        name: "get_solana_token_balance_by_mint_address",
+        description: "This {{tool}} gets the balance of an SPL token by its mint address",
+        parameters: getTokenBalanceByMintAddressParametersSchema,
+        method: async (
+            walletClient: SolanaWalletClient,
+            parameters: z.infer<typeof getTokenBalanceByMintAddressParametersSchema>,
+        ) => balanceOf(connection, parameters.wallet, parameters.mintAddress),
+    });
+
     return tools;
 }