From 61220ceb09caebdd405797d1b6f3454f0d6b3c1e Mon Sep 17 00:00:00 2001
From: Alejandro-Morales <alejandro.morales.mad@gmail.com>
Date: Fri, 29 Sep 2023 08:18:21 -0600
Subject: [PATCH] feat: added agent ledger and balance properties

---
 .../examples/13-agent-name-service/agent1.py  | 10 +++----
 python/src/uagents/agent.py                   | 26 +++++++++++++++++--
 python/src/uagents/config.py                  |  2 +-
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/python/examples/13-agent-name-service/agent1.py b/python/examples/13-agent-name-service/agent1.py
index c6abf1ba..113a1b92 100644
--- a/python/examples/13-agent-name-service/agent1.py
+++ b/python/examples/13-agent-name-service/agent1.py
@@ -1,6 +1,6 @@
 from cosmpy.aerial.wallet import LocalWallet
 
-from uagents.network import get_faucet, get_name_service_contract
+from uagents.network import get_faucet, get_name_service_contract, _testnet_ledger
 from uagents.setup import fund_agent_if_low
 from uagents import Agent, Context, Model
 
@@ -9,8 +9,6 @@
 
 # NOTE: Run agent1.py before running agent2.py
 
-# pylint: disable=protected-access
-
 
 class Message(Model):
     message: str
@@ -30,9 +28,9 @@ class Message(Model):
 DOMAIN = "agent"
 
 faucet = get_faucet()
-AGENT_BALANCE = bob._ledger.query_bank_balance(my_wallet)
+WALLET_BALANCE = _testnet_ledger.query_bank_balance(my_wallet)
 
-if AGENT_BALANCE < REGISTRATION_FEE:
+if WALLET_BALANCE < REGISTRATION_FEE:
     print("Adding funds to wallet...")
     faucet.get_wealth(my_wallet)
     print("Adding funds to wallet...complete")
@@ -41,7 +39,7 @@ class Message(Model):
 @bob.on_event("startup")
 async def register_agent_name(ctx: Context):
     await name_service_contract.register(
-        bob._ledger, my_wallet, ctx.address[-65:], ctx.name, DOMAIN
+        bob.ledger, my_wallet, ctx.address[-65:], ctx.name, DOMAIN
     )
 
 
diff --git a/python/src/uagents/agent.py b/python/src/uagents/agent.py
index 41d08e68..0a0a8563 100644
--- a/python/src/uagents/agent.py
+++ b/python/src/uagents/agent.py
@@ -9,6 +9,7 @@
 
 from cosmpy.aerial.wallet import LocalWallet, PrivateKey
 from cosmpy.crypto.address import Address
+from cosmpy.aerial.client import LedgerClient
 
 from uagents.asgi import ASGIServer
 from uagents.context import (
@@ -316,6 +317,16 @@ def wallet(self) -> LocalWallet:
         """
         return self._wallet
 
+    @property
+    def ledger(self) -> LedgerClient:
+        """
+        Get the ledger of the agent.
+
+        Returns:
+            LedgerClient: The agent's ledger
+        """
+        return self._ledger
+
     @property
     def storage(self) -> KeyValueStore:
         """
@@ -356,6 +367,17 @@ def mailbox_client(self) -> MailboxClient:
         """
         return self._mailbox_client
 
+    @property
+    def balance(self) -> int:
+        """
+        Get the balance of the agent.
+
+        Returns:
+            int: Bank balance.
+        """
+
+        return self.ledger.query_bank_balance(Address(self.wallet.address()))
+
     @mailbox.setter
     def mailbox(self, config: Union[str, Dict[str, str]]):
         """
@@ -471,7 +493,7 @@ async def register(self):
             or list(self.protocols.keys())
             != self._almanac_contract.get_protocols(self._identity.address)
         ):
-            agent_balance = self._ledger.query_bank_balance(
+            agent_balance = self.ledger.query_bank_balance(
                 Address(self.wallet.address())
             )
 
@@ -484,7 +506,7 @@ async def register(self):
             self._logger.info("Registering on almanac contract...")
             signature = self.sign_registration()
             await self._almanac_contract.register(
-                self._ledger,
+                self.ledger,
                 self.wallet,
                 self._identity.address,
                 list(self.protocols.keys()),
diff --git a/python/src/uagents/config.py b/python/src/uagents/config.py
index 6383de00..c85999fb 100644
--- a/python/src/uagents/config.py
+++ b/python/src/uagents/config.py
@@ -1,7 +1,7 @@
 import logging
 import sys
 
-# from enum import Enum
+
 from typing import Any, Dict, List, Optional, Union
 
 from uvicorn.logging import DefaultFormatter