diff --git a/js/src/example.ts b/js/src/example.ts
index 52b302a..4dec7ae 100644
--- a/js/src/example.ts
+++ b/js/src/example.ts
@@ -4,11 +4,10 @@ import { Numberu64, generateRandomSeed } from './utils';
 import { CreateSchedule } from './state';
 import {
   create,
-  VESTING_PROGRAM_ID,
-  DEVNET_VESTING_PROGRAM_ID,
-  TOKEN_MINT,
   initializeUnlock,
   unlock,
+  getTokenMint,
+  getProgramId,
 } from './main';
 import { signAndSendInstructions } from '@bonfida/utils';
 
@@ -39,8 +38,6 @@ const LOCKED_AMOUNT = 10;
 
 /** Your RPC connection */
 const connection = new Connection('');
-const DEVNET = true;
-const program = DEVNET ? DEVNET_VESTING_PROGRAM_ID : VESTING_PROGRAM_ID;
 
 /** Do some checks before sending the tokens */
 const checks = async () => {
@@ -50,7 +47,7 @@ const checks = async () => {
 
   // @ts-ignore
   const parsed = tokenInfo.value.data.parsed;
-  if (parsed.info.mint !== TOKEN_MINT.toBase58()) {
+  if (parsed.info.mint !== getTokenMint(connection).toBase58()) {
     throw new Error('Invalid mint');
   }
   if (parsed.info.owner !== LOCK_OWNER.toBase58()) {
@@ -77,7 +74,7 @@ const lock = async () => {
 
   const instruction = await create(
     connection,
-    program,
+    getProgramId(connection),
     // @ts-ignore
     Buffer.from(seed),
     wallet.publicKey,
@@ -95,7 +92,7 @@ const initUnlock = async () => {
 
   const instruction = await initializeUnlock(
     connection,
-    VESTING_PROGRAM_ID,
+    getProgramId(connection),
     // @ts-ignore
     Buffer.from(LOCK_SEED),
   );
@@ -110,7 +107,7 @@ const withdraw = async () => {
 
   const instruction = await unlock(
     connection,
-    program,
+    getProgramId(connection),
     // @ts-ignore
     Buffer.from(LOCK_SEED),
   );
diff --git a/js/src/main.ts b/js/src/main.ts
index b322c8c..55f9aab 100644
--- a/js/src/main.ts
+++ b/js/src/main.ts
@@ -254,3 +254,14 @@ export function getProgramId(connection: Connection): PublicKey {
     return DEVNET_VESTING_PROGRAM_ID;
   } else return VESTING_PROGRAM_ID;
 }
+
+/**
+ * This function can be used to retrieve the token mint based on the connection
+ * @param connection The Solana RPC connection object
+ * @returns A PublicKey object representing the token mint
+ */
+export function getTokenMint(connection: Connection): PublicKey {
+  if (isDevnetConnection(connection)) {
+    return DEVNET_TOKEN_MINT;
+  } else return TOKEN_MINT;
+}