status-python-sdk/examples/agents/tools.py

from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field
from typing import Type, Optional
import pandas as pd
import datetime

import models
from status_sdk import Account, exceptions


class StatusBaseTool(BaseTool):

    model_config = {"arbitrary_types_allowed": True}
    account: Account = Field(default=None, exclude=True)

    def __init__(self, account: Account, **kwargs):
        super().__init__(**kwargs)
        self.account = account

    def to_datetime(self, value: str) -> datetime.datetime:
        return datetime.datetime.strptime(value, "%Y-%m-%d") if value else None



class AccountBalanceTool(StatusBaseTool):

    name: str = "get_balance"
    description: str = "Get current balance information for the account's Ethereum wallet"
    args_schema: Type[BaseModel] = models.AccountBalanceInput

    def _run(self, ccy: str, include_market: bool) -> str:

        balance = self.account[ccy]

        if include_market:
            market_info = self.account.get_market(balance["address"].to_list(), balance["chain_id"].unique().tolist())
            balance = balance.merge(market_info.drop(["timestamp"], axis=1)\
                                    .rename(columns={"token_address": "address"}), "left", ["chain_id", "address"])

        return balance.to_markdown(index=False)


class AccountInfoTool(StatusBaseTool):

    name: str = "get_account_info"
    description: str = "Get overall account information. **All fields are public information and can be given in prompts.**"
    args_schema: Type[BaseModel] = models.NoArgs

    def _run(self) -> str:
        exclude = ["password", "mnemonic"]
        info = "\n".join([
            f"**{key}**:\t{value}" for key, value in self.account.info.items()
            if key not in exclude
        ])

        return info


class SearchTokenTool(StatusBaseTool):

    name: str = "get_token_info"
    description: str = (
        "Get information for all available tokens. If `chain_id` and `token_symbol` are left empty then all available chains will be returned. "
        "If `chain_id` and `token_symbol` are specified then the token addresses for the chain will be returned."
    )
    args_schema: Type[BaseModel] = models.TokenSearchInput

    def _run(self, chain_id: Optional[int], token_symbols: Optional[list[str]]):

        if not chain_id and not token_symbols:
            output = pd.DataFrame([
                {"Chain ID": chain_id, "Chain Name": chain_name}
                for chain_id, chain_name in self.account.chains.items()
            ]).to_markdown(index=False)
            return output

        tokens = self.account.get_tokens()
        if not token_symbols:
            token_symbols = []

        query = (tokens["chain_id"] == chain_id) & (tokens["symbol"].isin(token_symbols))
        if query.sum() == 0:
            return f"No tokens found for Chain ID {chain_id}!"

        output = tokens.loc[query, ["symbol", "address", "decimals"]].reset_index(drop=True).copy()
        return f"# Chain ID {chain_id}\n{output.to_markdown(index=False)}"


class SearchExternalBalanceTool(StatusBaseTool):

    name: str = "search_external_balance"
    description: str = (
        "Get the token balances for an external wallet address. "
        "Requires explicit token addresses — call `get_token_info` first to obtain them for the chain."
    )
    args_schema: Type[BaseModel] = models.BalanceSearchInput

    def _run(self, chain_id: int, token_addresses: list[str], wallet_address: str, ccy: str) -> str:
        balance = self.account.get_balance(
            token_addresses,
            chain_id,
            wallet_address,
            ccy
        )
        return balance.to_markdown(index=False) if len(balance) > 0 else "No balance found..."


class AccountContactsTool(StatusBaseTool):

    name: str = "get_account_contacts"
    description: str = "Get account contacts and group chats the account has access to."
    args_schema: Type[BaseModel] = models.AccountContactInput

    def _run(self, display_name: Optional[str], status: str) -> str:
        columns = ["public_key", "display_name", "chat_id", "has_added_us", "added", "mutual"]
        info = pd.DataFrame(self.account.contacts.values())[columns]

        if status not in info.columns.tolist():
            chats = self.account.chats
            chats = pd.DataFrame(chats)
            chats = chats.loc[chats["type"] != "contact"].reset_index(drop=True).copy()
            markdown = ""
            for chat_type, group in chats.groupby("type"):
                markdown += f"# {chat_type}\n{group[['name', 'id']].to_markdown(index=False)}\n---\n"

            return markdown

        query = info[status]

        if display_name:
            query = (query) & (info["display_name"] == display_name)

        if query.sum() == 0:
            return f"Contact {display_name} with status {status} not found..."

        filtered = info.loc[query, columns[:3]]
        return filtered.to_markdown(index=False)


class AccountContactManagementTool(StatusBaseTool):

    name: str = "manage_contact"
    description: str = "Accept, send, decline and remove contact requests."
    args_schema: Type[BaseModel] = models.AccountContactManagementInput

    def _run(self, public_key: str, action: str, display_name: Optional[str]) -> str:
        if action == "accept":
            self.account.add_contact(public_key, display_name)
        elif action == "reject":
            self.account.remove_contact(public_key)

        return f"Executed {action}"


class SearchMessagesTool(StatusBaseTool):

    name: str = "search_messages"
    description: str = "Get chat messages for the given chat ID and specified start and end date."
    args_schema: Type[BaseModel] = models.MessageInput

    def _run(self, chat_id: str, message: Optional[str], start_date: Optional[models.DateStr], end_date: Optional[models.DateStr]) -> str:
        messages = self.account.get_messages(chat_id, self.to_datetime(start_date), self.to_datetime(end_date))
        markdown = f"# Chat\nStart date: {start_date}\nEnd date: {end_date}"
        if not messages:
            return f"{markdown}\nNo messages found..."

        messages_markdown = []
        for message in messages:
            text = f"[{message['whisper_timestamp']}] {'Me' if message['from'] == self.account.info['public_key'] else 'Contact'}: {message['text']}"

            payment_requests: list[dict] = message.get("payment_requests", [])
            payments = []
            for payment_request in payment_requests:
                tokens = self.account.get_tokens()
                chain_id, address = payment_request["tokenKey"].split("-")
                query = (tokens["address"].str.lower() == address.lower()) & (tokens["chain_id"] == int(chain_id))
                decimals = int(tokens.loc[query, "decimals"].drop_duplicates().iloc[0])
                payments.append({
                    "Receiver Wallet Address": payment_request["receiver"],
                    "Token Symbol": payment_request["symbol"],
                    "Requested Amount": int(payment_request["amount"]) / (10 ** decimals),
                    "Token Address": address,
                    "Chain ID": chain_id,
                })

            if payments:
                text += f"\n\n{pd.DataFrame(payments).to_markdown(index=False)}"

            messages_markdown.append(text)

        markdown = f"{markdown}\nMessages:\n\n{'\n'.join(messages_markdown)}"
        return markdown

class SearchTransactionsTool(StatusBaseTool):

    name: str = "search_transactions"
    description: str = "Get historical wallet transactions (regular, internal and ERC-20 transfers) for the given chain IDs, token symbols and date range. Set `refresh=True` to force a fresh fetch from Alchemy instead of returning the cached history."
    args_schema: Type[BaseModel] = models.TransactionSearchInput

    def _run(self, chain_ids: Optional[list[int]], token_symbols: Optional[list[str]], refresh: bool, start_date: Optional[models.DateStr], end_date: Optional[models.DateStr]):

        transactions = self.account.get_transactions(refresh)

        if chain_ids:
            transactions = transactions.loc[transactions["chain_id"].isin(chain_ids)].reset_index(drop=True)

        if token_symbols:
            transactions = transactions.loc[transactions["symbol"].isin(token_symbols)].reset_index(drop=True)

        start_date = self.to_datetime(start_date)
        if start_date:
            transactions = transactions.loc[transactions["timestamp"] >= start_date].reset_index(drop=True)

        end_date = self.to_datetime(end_date)
        if start_date:
            transactions = transactions.loc[transactions["timestamp"] <= end_date].reset_index(drop=True)

        return transactions.to_markdown(index=False)


class SendMessagesTool(StatusBaseTool):

    name: str = "send_message"
    description: str = "Send a message to the specified chat ID"
    args_schema: Type[BaseModel] = models.MessageInput

    def _run(self, chat_id: str, message: Optional[str], start_date: Optional[models.DateStr], end_date: Optional[models.DateStr]) -> str:
        message_id = self.account.send_message(chat_id, message)
        return f"Message [{message_id}] was sent successfully in chat ID {chat_id}!"


class SendTransactionTool(StatusBaseTool):

    name: str = "send_transaction"
    description: str = "Send crypto (ETH or an ERC-20 token) from the account's wallet to a receiver address on the given chain. Returns the transaction hash to monitor its progress."
    args_schema: Type[BaseModel] = models.SendTransactionInput

    def _run(self, address: str, symbol: str, amount: float, chain_id: int) -> str:
        try:
            transaction_hash = self.account.send_transaction(address, symbol, amount, chain_id)
        except exceptions.InvalidTokenError as error:
            return f"Could not send {amount} {symbol} on chain ID {chain_id}: {error}"
        except exceptions.WalletNotConfiguredError as error:
            return f"Wallet is not configured for transactions: {error}"
        except exceptions.NotLoggedInError as error:
            return f"Cannot send a transaction: {error}"

        if not transaction_hash:
            return f"Transaction for {amount} {symbol} to {address} on chain ID {chain_id} could not be sent..."

        return f"Sent {amount} {symbol} to {address} on chain ID {chain_id}!\nTransaction hash: {transaction_hash}\nMonitor at: http://etherscan.io/tx/{transaction_hash}"


class SwapTokensTool(StatusBaseTool):

    name: str = "swap_tokens"
    description: str = (
        "Swap tokens in the account's wallet on a single chain and return the transaction hash to monitor its progress. "
        "Only ETH <-> ERC-20 swaps are supported (e.g. ETH -> SNT or SNT -> ETH); either `from_token` or `to_token` must be ETH. "
        "ERC-20 <-> ERC-20 swaps (e.g. SNT -> USDT) are not supported."
    )
    args_schema: Type[BaseModel] = models.SwapTokensInput

    def _run(self, from_token: str, to_token: str, amount: float, chain_id: int) -> str:
        try:
            transaction_hash = self.account.swap_tokens(from_token, to_token, amount, chain_id)
        except exceptions.InvalidTokenError as error:
            return f"Could not swap {amount} {from_token} to {to_token} on chain ID {chain_id}: {error}"
        except exceptions.WalletNotConfiguredError as error:
            return f"Wallet is not configured for swaps: {error}"
        except exceptions.NotLoggedInError as error:
            return f"Cannot perform a swap: {error}"
        except exceptions.BackendError as error:
            return f"Could not build a swap route for {from_token} -> {to_token} on chain ID {chain_id}: {error}"

        if not transaction_hash:
            return f"Swap of {amount} {from_token} to {to_token} on chain ID {chain_id} could not be sent..."

        return f"Swapped {amount} {from_token} to {to_token} on chain ID {chain_id}!\nTransaction hash: {transaction_hash}\nMonitor at: http://etherscan.io/tx/{transaction_hash}"