Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Get Balance for Other Wallets Fixes #291 #293

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+17 −8
Open

Get Balance for Other Wallets Fixes #291 #293

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1298d04
changed logic to support both current wallet + other wallets which us…
arnavmehta7 Feb 7, 2025
5758032
cleaned logic
arnavmehta7 Feb 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 17 additions & 8 deletions python/cdp-agentkit-core/cdp_agentkit_core/actions/get_balance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
from collections.abc import Callable
from typing import Optional

from cdp import Wallet
from cdp import Cdp, Wallet
from pydantic import BaseModel, Field

from cdp_agentkit_core.actions import CdpAction

GET_BALANCE_PROMPT = """
This tool will get the balance of all the addresses in the wallet for a given asset.
This tool will get the balance of address.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please see suggestion

Suggested change
This tool will get the balance of address.
This tool will get the balance of all the addresses in the wallet for a given asset, or for a user-specified address if provided.

It takes the asset ID as input. Always use 'eth' for the native asset ETH and 'usdc' for USDC.
"""

Expand All @@ -18,10 +19,14 @@ class GetBalanceInput(BaseModel):
...,
description="The asset ID to get the balance for, e.g. `eth`, `0x036CbD53842c5426634e7929541eC2318f3dCF7e`",
)
address_to_fund: Optional[str] = Field(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: rename to just address

None,
description="The address to fund. If not provided, the default address of the wallet will be used.",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please see suggestion

Suggested change
description="The address to fund. If not provided, the default address of the wallet will be used.",
description="The address to get the balance for. If not provided, addresses of the wallet will be used.",

)


def get_balance(wallet: Wallet, asset_id: str) -> str:
"""Get balance for all addresses in the wallet for a given asset.
def get_balance(wallet: Wallet, asset_id: str, address_to_fund: Optional[str]) -> str:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: rename to just address

"""Get balance for the address defined

Args:
wallet (Wallet): The wallet to get the balance for.
Expand All @@ -33,11 +38,15 @@ def get_balance(wallet: Wallet, asset_id: str) -> str:
"""
# for each address in the wallet, get the balance for the asset
balances = {}

try:
for address in wallet.addresses:
balance = address.balance(asset_id)
balances[address.address_id] = balance
if address_to_fund:
balances[address_to_fund] = Cdp.api_clients.external_addresses.get_external_address_balance(
network_id=wallet.network_id, address_id=address_to_fund, asset_id=asset_id
)
else:
for address in wallet.addresses:
balance = address.balance(asset_id)
balances[address.address_id] = balance
except Exception as e:
return f"Error getting balance for all addresses in the wallet {e!s}"

Expand Down
Loading