diff --git a/src/agent/index.ts b/src/agent/index.ts
index 90e47ef..d8086b8 100644
--- a/src/agent/index.ts
+++ b/src/agent/index.ts
@@ -84,7 +84,7 @@ export class SolanaAgentKit {
     return deploy_collection(this, options);
   }
 
-  async getBalance(token_address?: PublicKey): Promise<number | null> {
+  async getBalance(token_address?: PublicKey): Promise<number> {
     return get_balance(this, token_address);
   }
 
diff --git a/src/tools/get_balance.ts b/src/tools/get_balance.ts
index 36cdfef..05a56a1 100644
--- a/src/tools/get_balance.ts
+++ b/src/tools/get_balance.ts
@@ -10,7 +10,7 @@ import { SolanaAgentKit } from "../index";
 export async function get_balance(
   agent: SolanaAgentKit,
   token_address?: PublicKey,
-): Promise<number | null> {
+): Promise<number> {
   if (!token_address)
     return (
       (await agent.connection.getBalance(agent.wallet_address)) /
@@ -19,5 +19,5 @@ export async function get_balance(
 
   const token_account =
     await agent.connection.getTokenAccountBalance(token_address);
-  return token_account.value.uiAmount;
+  return token_account.value.uiAmount || 0;
 }