src/config.py

#!/usr/bin/env python3
"""
Shared configuration.

Includes:
1. System prompt templates
2. Token mappings (symbol -> contract address)
3. Tool definitions
4. Supported chains
"""

# ============================================================
# Token mappings
# ============================================================

# Tokens on Solana
SOLANA_TOKENS = {
    # Native
    "SOL": "So11111111111111111111111111111111111111112",
    
    # Stablecoins
    "USDC": "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
    "USDT": "Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB",
    
    # Major tokens
    "RAY": "4k3Dyjzvzp8eMZWUXbBCjEvwSkkk59S5iCNLY3QrkX6R",
    "JUP": "JUPyiwrYJFskUPiHa7hkeR8VUtAeFoSYbKedZNsDvCN",
    "BONK": "DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263",
    "WIF": "EKpQGSJtjMFqKZ9KQanSqYXRcF8fBopzLHYxdM65zcjm",
    "ORCA": "orcaEKTdK7LKz57vaAYr9QeNsVEPfiu6QeMU1kektZE",
    "MNGO": "MangoCzJ36AjZyKwVj3VnYU4GTonjfVEnJmvvWaxLac",
    "SRM": "SRMuApVNdxXokk5GT7XD5cUUgXMBCoAz2LHeuAoKWRt",
    "STEP": "StepAscQoEioFxxWGnh2sLBDFp9d8rvKz2Yp39iDpyT",
    "COPE": "8HGyAAB1yoM1ttS7pXjHMa3dukTFGQggnFFH3hJZgzQh",
    "FIDA": "EchesyfXePKdLtoiZSL8pBe8Myagyy8ZRqsACNCFGnvp",
    "MAPS": "MAPS41MDahZ9QdKXhVa4dWB9RuyfV4XqhyAZ8XcYepb",
    "OXY": "z3dn17yLaGMKffVogeFHQ9zWVcXgqgf3PQnDsNs2g6M",
    "MEDIA": "ETAtLmCmsoiEEKfNrHKJ2kYy3MoABhU6NQvpSfij5tDs",
    "SLND": "SLNDpmoWTVADgEdndyvWzroNL7zSi1dF9PC3xHGtPwp",
    "SAMO": "7xKXtg2CW87d97TXJSDpbD5jBkheTqA83TZRuJosgAsU",
    
    # Meme tokens
    "POPCAT": "7GCihgDB8fe6KNjn2MYtkzZcRjQy3t9GHdC8uHYmW2hr",
    "TRUMP": "6p6xgHyF7AeE6TZkSmFsko444wqoP15icUSqi2jfGiPN",
    "MELANIA": "FUAfBo2jgks6gB4Z4LfZkqSZgzNucisEHqnNebaRxM1P",
    "PEPE": "CLoUDKc4Ane7HeQcPpE3YHnznRxhMimJ4MyaUqyHFzAu",
    "DOGE": "9nEqaUcb16sQ3Tn1psbkWqyhPdLmfHWjKGymREjsAgTE",
    "SHIB": "CiKu4eHsVrc1eueVQeHn7qhXTcVu95gSQmBpX4utjL9z",
    "FLOKI": "FLKiaqMsaLwSZLqnCpSFHktMUQsXgcrhFN3TZg1Zj2EN",
    "WOJAK": "8cMrxCkREwszByAj9hzn5qKLnLjvLDU7gPYtWQgfwSAS",
    "GIGA": "63LfDmNb3MQ8mw9MtZ2To9bEA2M71kZUUGq5tiJxcqj9",
    "BOME": "ukHH6c7mMyiWCf1b9pnWe25TSpkDDt3H5pQZgZ74J82",
    "SLERF": "7BgBvyjrZX1YKz4oh9mjb8ZScatkkwb8DzFx7LoiVkM3",
    "MEW": "MEW1gQWJ3nEXg2qgERiKu7FAFj79PHvQVREQUzScPP5",
    "MYRO": "HhJpBhRRn4g56VsyLuT8DL5Bv31HkXqsrahTTUCZeZg4",
    "PONKE": "5z3EqYQo9HiCEs3R84RCDMu2n7anpDMxRhdK8PSWmrRC",
    "MOODENG": "ED5nyyWEzpPPiWimP8vYm7sD7TD3LAt3Q3gRTWHzPJBY",
    
    # DeFi tokens
    "PYTH": "HZ1JovNiVvGrGNiiYvEozEVgZ58xaU3RKwX8eACQBCt3",
    "JTO": "jtojtomepa8beP8AuQc6eXt5FriJwfFMwQx2v2f9mCL",
    "W": "85VBFQZC9TZkfaptBWjvUw7YbZjy52A6mjtPGjstQAmQ",
    "RENDER": "rndrizKT3MK1iimdxRdWabcF7Zg7AR5T4nud4EkHBof",
    "HNT": "hntyVP6YFm1Hg25TN9WGLqM12b8TQmcknKrdu1oxWux",
    "INJ": "6McPRfPV6bY1e9hLxWyG54W9i9Epq75QBvXg2oetBVTB",
    
    # Wrapped tokens
    "WETH": "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs",  # same as WETH
    "WBTC": "3NZ9JMVBmGAqocybic2c7LQCJScmgsAZ6vQqTDzcqmJh",  # same as WBTC
    "ETH": "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs",  # same as WETH
    "BTC": "3NZ9JMVBmGAqocybic2c7LQCJScmgsAZ6vQqTDzcqmJh",  # same as WBTC
}

# Tokens on Ethereum
ETHEREUM_TOKENS = {
    "ETH": "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE",
    "WETH": "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
    "USDC": "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
    "USDT": "0xdAC17F958D2ee523a2206206994597C13D831ec7",
    "DAI": "0x6B175474E89094C44Da98b954EescdeCB5BE3830",
    "WBTC": "0x2260FAC5E5542a773Aa44fBCfeDf7C193bc2C599",
    "UNI": "0x1f9840a85d5aF5bf1D1762F925BDADdC4201F984",
    "LINK": "0x514910771AF9Ca656af840dff83E8264EcF986CA",
    "AAVE": "0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9",
    "PEPE": "0x6982508145454Ce325dDbE47a25d4ec3d2311933",
    "SHIB": "0x95aD61b0a150d79219dCF64E1E6Cc01f0B64C4cE",
    "TRUMP": "0x576e2BeD8F7b46D34016198911Cdf9886f78bea7",
}

# Tokens on BSC
BSC_TOKENS = {
    "BNB": "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE",
    "WBNB": "0xbb4CdB9CBd36B01bD1cBaEBF2De08d9173bc095c",
    "USDC": "0x8AC76a51cc950d9822D68b83fE1Ad97B32Cd580d",
    "USDT": "0x55d398326f99059fF775485246999027B3197955",
    "BUSD": "0xe9e7CEA3DedcA5984780Bafc599bD69ADd087D56",
    "CAKE": "0x0E09FaBB73Bd3Ade0a17ECC321fD13a19e81cE82",
    "DOGE": "0xbA2aE424d960c26247Dd6c32edC70B295c744C43",
    "SHIB": "0x2859e4544C4bB03966803b044A93563Bd2D0DD4D",
    "FLOKI": "0xfb5B838b6cfEEdC2873aB27866079AC55363D37E",
}

# Tokens on Base
BASE_TOKENS = {
    "ETH": "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE",
    "WETH": "0x4200000000000000000000000000000000000006",
    "USDC": "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913",
    "DAI": "0x50c5725949A6F0c72E6C4a641F24049A917DB0Cb",
    "BRETT": "0x532f27101965dd16442E59d40670FaF5eBB142E4",
    "DEGEN": "0x4ed4E862860beD51a9570b96d89aF5E1B0Efefed",
    "TOSHI": "0xAC1Bd2486aAf3B5C0fc3Fd868558b082a531B2B4",
}

# Token mappings for all chains
ALL_TOKENS = {
    "solana": SOLANA_TOKENS,
    "ethereum": ETHEREUM_TOKENS,
    "bsc": BSC_TOKENS,
    "base": BASE_TOKENS,
}

# Supported chains
SUPPORTED_CHAINS = ["solana", "ethereum", "bsc", "base"]

# Default chain
DEFAULT_CHAIN = "solana"


# ============================================================
# System prompt templates
# ============================================================

def get_system_prompt(chain: str = "solana") -> str:
    """
    Build a system prompt that contains the token mapping for a chain.
    
    Args:
        chain: blockchain name
    
    Returns:
        system prompt string
    """
    tokens = ALL_TOKENS.get(chain, SOLANA_TOKENS)
    
    # Build the token mapping string
    token_list = []
    for symbol, address in tokens.items():
        token_list.append(f"  - {symbol}: {address}")
    token_mapping_str = "\n".join(token_list)
    
    system_prompt = f"""You are a blockchain trading assistant that helps users search for tokens and execute swaps.

## Available Tools

1. **SEARCH_TOKEN**: Search for token information on the blockchain.
   - Parameters: symbol, address, chain, keyword
   - Use when user wants to find/query/search token information

2. **EXECUTE_SWAP**: Execute a token swap on the blockchain.
   - Parameters: inputTokenCA, outputTokenCA, inputTokenAmount, inputTokenPercentage
   - Use when user wants to buy/sell/swap/convert tokens

## Token Address Mapping ({chain.upper()})

{token_mapping_str}

## Important Rules

1. When user mentions a token symbol (like SOL, USDC, RAY), use the corresponding contract address from the mapping above.
2. For buying tokens: inputTokenCA is what user spends (usually SOL/ETH), outputTokenCA is what user receives.
3. For selling tokens: inputTokenCA is what user sells, outputTokenCA is usually SOL/ETH.
4. If user specifies an amount, use inputTokenAmount. If user specifies a percentage (like "all", "50%", "half"), use inputTokenPercentage.
5. inputTokenPercentage should be a decimal between 0 and 1 (e.g., 0.5 for 50%, 1.0 for 100%).
6. If the user's request is unclear or missing required information, ask for clarification instead of making assumptions.
7. Do NOT call any function if the request is unrelated to token search or swap (e.g., weather, greetings, jokes).

## Examples

- "Buy 2 SOL of RAY" → EXECUTE_SWAP with inputTokenCA=SOL address, outputTokenCA=RAY address, inputTokenAmount="2"
- "Sell all my JUP" → EXECUTE_SWAP with inputTokenCA=JUP address, outputTokenCA=SOL address, inputTokenPercentage=1.0
- "Search for BONK" → SEARCH_TOKEN with symbol="BONK", chain="{chain}"
"""
    
    return system_prompt


# Shorter system prompt (used for training)
def get_system_prompt_short(chain: str = "solana") -> str:
    """Build the concise system prompt."""
    tokens = ALL_TOKENS.get(chain, SOLANA_TOKENS)
    
    # Only include common tokens
    common_tokens = ["SOL", "USDC", "USDT", "RAY", "JUP", "BONK", "WIF", "TRUMP", "POPCAT", "PEPE", "ETH", "WETH", "BTC", "WBTC"]
    token_list = []
    for symbol in common_tokens:
        if symbol in tokens:
            token_list.append(f"{symbol}:{tokens[symbol]}")
    token_mapping_str = ", ".join(token_list)
    
    system_prompt = f"""You are a blockchain trading assistant. Use SEARCH_TOKEN to find tokens and EXECUTE_SWAP to trade tokens.

Token addresses ({chain}): {token_mapping_str}

Rules:
- For buy: inputTokenCA=payment token, outputTokenCA=target token
- For sell: inputTokenCA=token to sell, outputTokenCA=SOL/ETH
- Use inputTokenAmount for specific amounts, inputTokenPercentage (0-1) for percentages
- Ask for clarification if request is unclear
- Do NOT call functions for unrelated requests (weather, greetings, etc.)
"""
    
    return system_prompt


# ============================================================
# Tool definitions
# ============================================================

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "SEARCH_TOKEN",
            "description": "Search for token information on the blockchain. Use this when user wants to find, query, or get information about a token.",
            "parameters": {
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": ["string", "null"],
                        "description": "Symbol of the token (e.g., SOL, USDC, RAY)"
                    },
                    "address": {
                        "type": ["string", "null"],
                        "description": "Contract address of the token"
                    },
                    "chain": {
                        "type": "string",
                        "enum": ["solana", "ethereum", "bsc", "base"],
                        "description": "Blockchain to search on"
                    },
                    "keyword": {
                        "type": ["string", "null"],
                        "description": "Keyword to search for the token"
                    }
                },
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "EXECUTE_SWAP",
            "description": "Execute a token swap on the blockchain. Use this when user wants to buy, sell, swap, or convert tokens.",
            "parameters": {
                "type": "object",
                "properties": {
                    "inputTokenCA": {
                        "type": ["string", "null"],
                        "description": "Contract address of the token to spend/sell"
                    },
                    "outputTokenCA": {
                        "type": ["string", "null"],
                        "description": "Contract address of the token to receive/buy"
                    },
                    "inputTokenAmount": {
                        "type": ["string", "null"],
                        "description": "Exact amount of input token to swap"
                    },
                    "inputTokenPercentage": {
                        "type": ["number", "null"],
                        "description": "Percentage of input token balance to swap (0-1, e.g., 0.5 for 50%)"
                    }
                },
                "required": ["inputTokenCA", "outputTokenCA"]
            }
        }
    }
]


# ============================================================
# Helper functions
# ============================================================

def get_token_address(symbol: str, chain: str = "solana") -> str:
    """Get token address by symbol."""
    tokens = ALL_TOKENS.get(chain, SOLANA_TOKENS)
    return tokens.get(symbol.upper(), None)


def get_native_token(chain: str = "solana") -> tuple:
    """Get the native token (symbol, address) for a chain."""
    native_tokens = {
        "solana": ("SOL", SOLANA_TOKENS["SOL"]),
        "ethereum": ("ETH", ETHEREUM_TOKENS["ETH"]),
        "bsc": ("BNB", BSC_TOKENS["BNB"]),
        "base": ("ETH", BASE_TOKENS["ETH"]),
    }
    return native_tokens.get(chain, ("SOL", SOLANA_TOKENS["SOL"]))


def get_all_token_symbols(chain: str = "solana") -> list:
    """Return all token symbols on a chain."""
    tokens = ALL_TOKENS.get(chain, SOLANA_TOKENS)
    return list(tokens.keys())


if __name__ == "__main__":
    # Quick self-test
    print("=== System Prompt (Full) ===")
    print(get_system_prompt("solana")[:500] + "...")
    print("\n=== System Prompt (Short) ===")
    print(get_system_prompt_short("solana"))
    print("\n=== Token Address ===")
    print(f"RAY: {get_token_address('RAY', 'solana')}")
    print(f"TRUMP: {get_token_address('TRUMP', 'solana')}")