Add the 'balance' function