pydantic_agent.py

# PydanticAI Agent with MCP

from pydantic_ai import Agent, RunContext
from pydantic_ai.mcp import MCPServerStdio
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

from dotenv import load_dotenv
import os


load_dotenv()

live = True
if live:
    env = {
        'ARMOR_API_KEY': os.getenv('ARMOR_API_KEY'),
    }
else:
    env = {
        'ARMOR_API_KEY': os.getenv('ARMOR_API_KEY'),
    }

# what are the top coins? figure out which one is the best short term (couple hours) investment from SOL
# now use 0.7 SOL to buy your best pick, and then set 20 DCA sell order for the optimal time window 100% back to SOL

# Set up OpenRouter based model
API_KEY = os.getenv('OPENROUTER_API_KEY')
model = OpenAIModel(
    'anthropic/claude-3.7-sonnet',
    provider=OpenAIProvider(
        base_url='https://openrouter.ai/api/v1', 
        api_key=API_KEY
    ),
)

# MCP setup for armor-crypto-mcp
# Toggle between first and second line to test different things
# Immediate test can be done by git cloning armor-crypto-mcp repo into this same folder

mcp_servers = [
    # MCPServerStdio('uvx', ['armor-crypto-mcp'], env=env),
    MCPServerStdio('python', ['armor-crypto-mcp/armor_crypto_mcp/armor_mcp.py'], env=env),
]

# first_message = "do a search for btcusd and tell me the current price of btc"
first_message = "List my Wallets"

# Set up Agent with Server
agent = Agent(model, mcp_servers=mcp_servers, system_prompt="""
You are a top crypto expert with access to the Armor suite of tools. You help the user to make transactions, strategize, swap and study the crypto ecosystem.

## Response Instructions
- Talk like WallStreetBets stonk ape. You ape stronk ape! crypto stonks gud! Help user get bananna 🦍🦍🦍
- Reply as if speaking, optimizing the output and simpifying complex data
- Apes together stronk!
- Assume the user wants to know what is in their wallets after it's contents may have changed, wait a bit after a transaction and get all the wallets, try to calculate pnl
- If the user is asking to trade something risky with a lot of money, verify the transaction, giving all of the details for review
- Try to convince the user away from being exit liquidity on any positions they may consider.

## Group Operations
- When asked to make a swap or stake with a group of wallets, use an equal amount of tokens from every wallet in that group
- Only use wallets which have enough to contribute evenly to the operation

## Predefined Tasks
- Get Banana
  - check wallets
  - look up top tokens
  - look up current SOL
  - do analysis to find which is best 1hr gain based data you find, obviously discarding shitcoins
  - show juciest banannas to user with a strategy on each one, let them choose next move

""")

async def main():
    """Simple toy loop for testing agent functionalities with prompts and basic memory"""
    async with agent.run_mcp_servers(): 
        result = await agent.run(first_message)
        user_input = ""
        message_history = result.new_messages()
        while True:
            print(f"\n{result.data}")
            user_input = input("\n> ")
            # print([type(v) for v in message_history])
            result = None
            err = None
            for i in range(0, 3):
                try:
                    result = await agent.run(
                        user_input, 
                        message_history=message_history
                    )
                    break
                except Exception as e:
                    err = e
                    await asyncio.sleep(2)
            if result is None:
                print(f"\nError {err}. Try again...\n")
                continue
            message_history.extend(result.new_messages())
            if len(message_history) > 12:
                message_history.pop(0)
        
if __name__ == "__main__":
    import asyncio
    asyncio.run(main())