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 bot import Account


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.strftime(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 balance for an external address."
    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 chad 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:
        to_datetime = lambda value: datetime.datetime.strptime(value, "%Y-%m-%d") if value else None
        messages = self.account.get_messages(chat_id, to_datetime(start_date), to_datetime(end_date))
        markdown = f"# Chat\nStart date: {start_date}\nEnd date: {end_date}"
        if messages:
            messages_markdown = [f"[{message['whisper_timestamp']}] {'Me' if message['from'] == self.account.info['public_key'] else 'Contact'}: {message['text']}" for message in messages]
            markdown = f"{markdown}\nMessages:\n{messages_markdown}"

        return markdown



class SendMessagesTool(StatusBaseTool):

    name: str = "send_message"
    description: str = "Send a message to the specified chad IT"
    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:
        self.account.send_message(chat_id, message)
        return f"Message was sent successfully in chat ID {chat_id}!"