treasuryManagement/portfolioAgent.py at main · janusdefi/treasuryManagement · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

from web3 import Web3
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()
INFURA_URL = os.getenv("INFURA_URL")
PRIVATE_KEY = os.getenv("PRIVATE_KEY")
WALLET_ADDRESS = os.getenv("WALLET_ADDRESS")
TREASURY_CONTRACT = os.getenv("TREASURY_CONTRACT")

# Connect to Ethereum node
web3 = Web3(Web3.HTTPProvider(INFURA_URL))
assert web3.isConnected(), "Failed to connect to Ethereum node"

# Treasury Contract ABI (Simplified for this example)
TREASURY_ABI = [
    {
        "inputs": [],
        "name": "getActiveAssets",
        "outputs": [{"internalType": "address[]", "name": "", "type": "address[]"}],
        "stateMutability": "view",
        "type": "function",
    },
    {
        "inputs": [],
        "name": "rebalancePortfolio",
        "outputs": [],
        "stateMutability": "nonpayable",
        "type": "function",
    },
]

# Load Treasury Contract
treasury_contract = web3.eth.contract(address=TREASURY_CONTRACT, abi=TREASURY_ABI)


# Fetch active assets
def get_active_assets():
    return treasury_contract.functions.getActiveAssets().call()


# Mock AI logic for rebalancing
def mock_ai_rebalance():
    active_assets = get_active_assets()
    print(f"Active Assets: {active_assets}")

    # Simulate rebalancing logic (assign equal weights for simplicity)
    new_allocations = [1000 for _ in active_assets]  # Mock allocations in token units
    return active_assets, new_allocations


# Submit rebalance transaction
def submit_rebalance(assets, allocations):
    # Build transaction
    nonce = web3.eth.getTransactionCount(WALLET_ADDRESS)
    gas_price = web3.eth.gas_price
    transaction = treasury_contract.functions.rebalancePortfolio().buildTransaction({
        "chainId": 5,  # Goerli Testnet
        "gas": 500000,
        "gasPrice": gas_price,
        "nonce": nonce,
        "from": WALLET_ADDRESS,
    })

    # Sign transaction
    signed_txn = web3.eth.account.sign_transaction(transaction, PRIVATE_KEY)

    # Send transaction
    txn_hash = web3.eth.send_raw_transaction(signed_txn.rawTransaction)
    print(f"Transaction sent! Hash: {txn_hash.hex()}")

    # Wait for confirmation
    receipt = web3.eth.wait_for_transaction_receipt(txn_hash)
    print(f"Transaction confirmed! Block: {receipt.blockNumber}")


# Main function to run the AI agent
def main():
    print("Fetching active assets...")
    assets, allocations = mock_ai_rebalance()

    print(f"Proposed Allocations: {allocations}")
    print("Submitting rebalance transaction...")
    submit_rebalance(assets, allocations)


if __name__ == "__main__":
    main()