mirror of
https://github.com/status-im/status-python-sdk.git
synced 2026-08-30 05:31:11 +00:00
280 lines
12 KiB
Python
280 lines
12 KiB
Python
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}"
|