Merge remote-tracking branch 'origin/main' into nerfzael/swap-tests

README.md

@@ -69,17 +69,26 @@ AutoTx can be connected to your existing smart account by doing the following:
 
 
 ## Prompts
-AutoTx currently supports prompts such as:
-* `Send 1 ETH to vitalik.eth`  
-* `Buy 100 USDC with ETH`  
-* `Swap ETH to 0.05 WBTC, then swap WBTC to 1000 USDC, and finally send 50 USDC to 0x...`  
+AutoTx currently supports prompts such as:  
+
+| Category | Prompt |
+|---|---|
+| Token Research | `Research the top AI coins by trading volume on CoinGecko` |
+| Token Research | `Conduct a thorough analysis of WLD, including whether to hold or sell` |
+| Token Research | `Find leveraged tokens I can buy directly on Ethereum mainnet`  |
+| Send Tokens | `Send tokens 10 DAI and 0.01 ETH to 0x...` |
+| Send Tokens, ENS | `Send tokens 1 ETH and 1000 USDC to vitalik.eth` |
+| Swap Tokens | `Buy 100 USDC with ETH` |
+| Swap Tokens | `Swap tokens directly through Uniswap` |
+| Multi Task | `Identify the top AI coins by trading volume on the Ethereum mainnet using CoinGecko. Buy 1 ETH of the top 2.` |
+| Multi Task| `Swap ETH to 0.05 WBTC, then swap WBTC to 1000 USDC, and finally send 50 USDC to 0x...` |
+| Multi Task, Airdrop | `Buy 10 WLD with ETH, then 50% of the purchased amount of WLD should be sent in equal amounts to each of these addresses: vitalik.eth, 0x...` |
 
 Future possibilities:
-* `Buy me 5 small cap AI tokens`
-* `Send the most popular meme coin to vitalik.eth`
 * `Purchase mainnet ETH with my USDC on optimism`
 * `What proposals are being voted on right now?`
-* `Donate $100 to environmental impact projects.`
+* `Donate $100 to environmental impact projects`
+* ...
 
 ## Use AutoTx With Open-Source Models
 
@@ -92,16 +101,21 @@ To run AutoTx with your favorite OS model, you can use any provider that simulat
 Now simply run AutoTx as normally do. For more tips on choosing the best model, you can follow [this guide](https://microsoft.github.io/autogen/docs/topics/non-openai-models/best-tips-for-nonopenai-models/).
 
 ## How To Contribute
-Interested in contributing to AutoTx? There's no shortage of [agents](#agents) to build! Additionally, checkout the [repository's issues](https://github.com/polywrap/AutoTx/issues) that will remain updated with the project's latest developments. Connect with us on [Discord](https://discord.gg/k7UCsH3ps9) if you have any questions or ideas to share.
+Interested in contributing to AutoTx? Here are some ideas:
+* Contribute prompt ideas above
+* Build an [agent](#agents)
+* Discuss AutoTx's future in [issues](https://github.com/polywrap/AutoTx/issues)
 
-### Adding Agents
+Connect with us on [Discord](https://discord.gg/k7UCsH3ps9) if you have any questions or ideas to share.
+
+## Building Agents
 
 To add agents to AutoTx, we recommend starting with the [`ExampleAgent.py`](./autotx/agents/ExampleAgent.py) starter template. From there you'll want to:
 1. Implement the tools (functions) you want the agent to be able to call.
 2. Add all tools to the agent's `tools=[...]` array.
 3. Add your new agent to `AutoTx`'s constructor in [`cli.py`](./autotx/cli.py).
 
-### Testing Agents
+### Testing
 
 Tests are located in the [`./autotx/tests`](./autotx/tests/) directory.
 

autotx/AutoTx.py

@@ -62,8 +62,10 @@ def run(self, prompt: str, non_interactive: bool, silent: bool = False) -> None:
                     Verifier analyzes chat and responds with TERMINATE if the goal is met.
                     Verifier can consider the goal met if the other agents have prepared the necessary transactions.
                     
-                    If some information needs to be returned to the user, add it in your answer and then say the word TERMINATE.
-                    Make sure to only add information if the user explicitly ask for a question that needs to be answered
+                    If some information needs to be returned to the user or if there are any errors encountered during the process, add this in your answer.
+                    Start any error messages with "ERROR:" to clearly indicate the issue. Then say the word TERMINATE.
+                    Make sure to only add information if the user explicitly asks for a question that needs to be answered
+                    or error details if user's request can not be completed.
                     """
                 ),
             llm_config=self.get_llm_config(),
@@ -98,7 +100,11 @@ def run(self, prompt: str, non_interactive: bool, silent: bool = False) -> None:
         ))
 
         if chat.summary:
-            cprint(chat.summary.replace("\n", ""), "green")
+            if "ERROR:" in chat.summary:
+                error_message = chat.summary.replace("ERROR: ", "").replace("\n", "")
+                cprint(error_message, "red")
+            else:
+                cprint(chat.summary.replace("\n", ""), "green")
 
         try:
             self.manager.send_tx_batch(self.transactions, require_approval=not non_interactive)

autotx/agents/SendTokensAgent.py

@@ -11,56 +11,28 @@
     get_erc20_balance,
 )
 from autotx.utils.ethereum import load_w3
-from autotx.utils.ethereum.build_transfer_eth import build_transfer_eth
+from autotx.utils.ethereum.build_transfer_native import build_transfer_native
 from web3.constants import ADDRESS_ZERO
+from autotx.utils.ethereum.constants import NATIVE_TOKEN_ADDRESS
 from autotx.utils.ethereum.eth_address import ETHAddress
+from autotx.utils.ethereum.get_native_balance import get_native_balance
 
 name = "send-tokens"
 
 system_message = lambda autotx: dedent(f"""
     You are an AI assistant. Assist the user (address: {autotx.manager.address}) in their tasks by fetching balances and preparing transactions to send tokens.
-    You are an expert in Ethereum tokens and can help users send tokens and check their balances.
+    You are an expert in Ethereum tokens (native and erc20) and can help users send tokens and check their balances.
     You use the tools available to assist the user in their tasks. 
     Your job is to only prepare the transactions and the user will take care of executing them.
-    NOTE: There is no reason to call get_erc20_balance after calling transfer as the transfers are only prepared and not executed. 
+    NOTE: There is no reason to call get_balance after calling transfer as the transfers are only prepared and not executed. 
     """
 )
 
-class TransferETHTool(AutoTxTool):
-    name: str = "prepare_transfer_eth_transaction"
+class TransferTokenTool(AutoTxTool):
+    name: str = "prepare_transfer_transaction"
     description: str = dedent(
         """
-        Prepares a transfer transaction for given amount in decimals for ETH.
-        """
-    )
-
-    def build_tool(self, autotx: AutoTx) -> Callable[[float, str], str]:
-        def run(
-            amount: Annotated[float, "Amount given by the user to transfer. The function will take care of converting the amount to needed decimals."],
-            receiver: Annotated[str, "The receiver's address or ENS domain"]
-        ) -> str:
-            web3 = load_w3()
-
-            receiver_addr = ETHAddress(receiver, web3)
-
-            tx = build_transfer_eth(web3, ETHAddress(ADDRESS_ZERO, web3), receiver_addr, amount)
-
-            prepared_tx = PreparedTx(f"Transfer {amount} ETH to {str(receiver_addr)}", tx)
-
-            autotx.transactions.append(prepared_tx)
-
-            print(f"Prepared transaction: {prepared_tx.summary}")
-
-            return prepared_tx.summary
-
-
-        return run
-
-class TransferERC20Tool(AutoTxTool):
-    name: str = "prepare_transfer_erc20_transaction"
-    description: str = dedent(
-        """
-        Prepares a transfer transaction for given amount in decimals for given token.
+        Prepares a transfer transaction for given amount in decimals for a token
         """
     )
 
@@ -70,52 +42,32 @@ def run(
             receiver: Annotated[str, "The receiver's address or ENS domain"],
             token: Annotated[str, "Symbol of token to transfer"]
         ) -> str:
-            token_address = Web3.to_checksum_address(autotx.network.tokens[token.lower()])
             web3 = load_w3()
-
             receiver_addr = ETHAddress(receiver, web3)
-
-            tx = build_transfer_erc20(web3, token_address, receiver_addr, amount)
-
+            token_address = ETHAddress(autotx.network.tokens[token.lower()], web3)
+
+            prepared_tx: PreparedTx | None = None
+
+            if token_address.hex == NATIVE_TOKEN_ADDRESS:
+                tx = build_transfer_native(web3, ETHAddress(ADDRESS_ZERO, web3), receiver_addr, amount)
+            else:
+                tx = build_transfer_erc20(web3, token_address, receiver_addr, amount)
+
             prepared_tx = PreparedTx(f"Transfer {amount} {token} to {str(receiver_addr)}", tx)
-        
+
             autotx.transactions.append(prepared_tx)
-
+            
             print(f"Prepared transaction: {prepared_tx.summary}")
-
+            
             return prepared_tx.summary
 
         return run
-
-class GetETHBalanceTool(AutoTxTool):
-    name: str = "get_eth_balance"
-    description: str = dedent(
-        """
-        Check owner balance in ETH
-        """
-    )
 
-    def build_tool(self, autotx: AutoTx) -> Callable[[str], float]:
-        def run(
-            owner: Annotated[str, "The owner's address or ENS domain"]
-        ) -> float:
-            web3 = load_w3()
-
-            owner_addr = ETHAddress(owner, web3)
-
-            eth_balance = web3.eth.get_balance(owner_addr.hex)
-
-            print(f"Fetching ETH balance for {str(owner_addr)}: {eth_balance / 10 ** 18} ETH")
-
-            return eth_balance / 10 ** 18
-
-        return run
-
-class GetERC20BalanceTool(AutoTxTool):
-    name: str = "get_erc20_balance"
+class GetTokenBalanceTool(AutoTxTool):
+    name: str = "get_token_balance"
     description: str = dedent(
         """
-        Check owner balance in ERC20 token
+        Check owner balance of token
         """
     )
 
@@ -125,23 +77,26 @@ def run(
             owner: Annotated[str, "The token owner's address or ENS domain"]
         ) -> float:
             web3 = load_w3()
-            token_address = ETHAddress(autotx.network.tokens[token.lower()], web3)
             owner_addr = ETHAddress(owner, web3)
+            token_address = ETHAddress(autotx.network.tokens[token.lower()], web3)
+
+            balance: float = 0
 
-            balance = get_erc20_balance(web3, token_address, owner_addr)
+            if token_address.hex == NATIVE_TOKEN_ADDRESS:
+                balance = get_native_balance(web3, owner_addr)
+            else:
+                balance = get_erc20_balance(web3, token_address, owner_addr)
 
             print(f"Fetching {token} balance for {str(owner_addr)}: {balance} {token}")
 
             return balance
 
         return run
-
+    
 class SendTokensAgent(AutoTxAgent):
     name = name
     system_message = system_message
     tools = [
-        TransferETHTool(),
-        TransferERC20Tool(),
-        GetETHBalanceTool(),
-        GetERC20BalanceTool()
+        TransferTokenTool(),
+        GetTokenBalanceTool(),
     ]

autotx/agents/SwapTokensAgent.py

@@ -72,14 +72,14 @@ def get_tokens_address(token_in: str, token_out: str, network_info: NetworkInfo)
 
     if not network_info.chain_id in SUPPORTED_UNISWAP_V3_NETWORKS:
         raise ChainIdNotSupported(
-            f"Network {network_info.chain_id.name} not supported for swap"
+            f"Network {network_info.chain_id.name.lower()} not supported for swap"
         )
 
     if token_in not in network_info.tokens:
-        raise Exception(f"Token {token_in} is not supported")
+        raise Exception(f"Token {token_in} is not supported in network {network_info.chain_id.name.lower()}")
 
     if token_out not in network_info.tokens:
-        raise Exception(f"Token {token_out} is not supported")
+        raise Exception(f"Token {token_out} is not supported in network {network_info.chain_id.name.lower()}")
 
     return (network_info.tokens[token_in], network_info.tokens[token_out])
 

autotx/cli.py

@@ -7,15 +7,15 @@
 from autotx.agents.ResearchTokensAgent import ResearchTokensAgent
 from autotx.agents.SendTokensAgent import SendTokensAgent
 from autotx.agents.SwapTokensAgent import SwapTokensAgent
-from autotx.utils.ethereum import get_eth_balance
+from autotx.utils.ethereum import get_native_balance
 
 from autotx.utils.constants import COINGECKO_API_KEY, OPENAI_API_KEY, OPENAI_MODEL_NAME
 from autotx.utils.ethereum.networks import NetworkInfo
 from autotx.utils.ethereum.helpers.get_dev_account import get_dev_account
 from autotx.AutoTx import AutoTx
 from autotx.utils.ethereum.agent_account import get_or_create_agent_account
 from autotx.utils.ethereum.SafeManager import SafeManager
-from autotx.utils.ethereum.send_eth import send_eth
+from autotx.utils.ethereum.send_native import send_native
 from autotx.utils.ethereum.helpers.show_address_balances import show_address_balances
 from autotx.utils.configuration import get_configuration
 
@@ -60,8 +60,8 @@ def run(prompt: str | None, non_interactive: bool, verbose: bool) -> None:
         manager = SafeManager.deploy_safe(client, dev_account, agent, [dev_account.address, agent.address], 1)
         print(f"Smart account deployed: {manager.address}")
 
-        if get_eth_balance(web3, manager.address) < 10:
-            send_eth(dev_account, manager.address, 10, web3)
+        if get_native_balance(web3, manager.address) < 10:
+            send_native(dev_account, manager.address, 10, web3)
             print(f"Sent 10 ETH to smart account for testing purposes")
 
         print("=" * 50)

autotx/tests/agents/token/send/test_send.py

@@ -1,16 +1,15 @@
 from autotx.utils.ethereum import get_erc20_balance
-from autotx.utils.ethereum.get_eth_balance import get_eth_balance
+from autotx.utils.ethereum.get_native_balance import get_native_balance
 
-
-def test_auto_tx_send_eth(configuration, auto_tx, test_accounts):
+def test_auto_tx_send_native(configuration, auto_tx, test_accounts):
     (_, _, client, _) = configuration
     receiver = test_accounts[0]
-    balance = get_eth_balance(client.w3, receiver)
+    balance = get_native_balance(client.w3, receiver)
     assert balance == 0
 
     auto_tx.run(f"Send 1 ETH to {receiver}", non_interactive=True)
 
-    balance = get_eth_balance(client.w3, receiver)
+    balance = get_native_balance(client.w3, receiver)
     assert balance == 1
 
 def test_auto_tx_send_erc20(configuration, auto_tx, usdc, test_accounts):
@@ -28,18 +27,18 @@ def test_auto_tx_send_erc20(configuration, auto_tx, usdc, test_accounts):
 
     assert balance + 10 == new_balance
 
-def test_auto_tx_send_eth_sequential(configuration, auto_tx, test_accounts):
+def test_auto_tx_send_native_sequential(configuration, auto_tx, test_accounts):
     (_, _, client, _) = configuration
     receiver = test_accounts[0]
 
     auto_tx.run(f"Send 1 ETH to {receiver}", non_interactive=True)
 
-    balance = get_eth_balance(client.w3, receiver)
+    balance = get_native_balance(client.w3, receiver)
     assert balance == 1
 
     auto_tx.run(f"Send 0.5 ETH to {receiver}", non_interactive=True)
 
-    balance = get_eth_balance(client.w3, receiver)
+    balance = get_native_balance(client.w3, receiver)
     assert balance == 1.5
 
 def test_auto_tx_send_erc20_parallel(configuration, auto_tx, usdc, test_accounts):
@@ -59,4 +58,35 @@ def test_auto_tx_send_erc20_parallel(configuration, auto_tx, usdc, test_accounts
     new_balance_two = get_erc20_balance(client.w3, usdc, receiver_two)
 
     assert balance_one + 2 == new_balance_one
-    assert balance_two + 3 == new_balance_two
+    assert balance_two + 3 == new_balance_two
+
+def test_auto_tx_send_eth_multiple(configuration, auto_tx, usdc, test_accounts):
+    (_, _, client, _) = configuration
+
+    receiver_1 = test_accounts[0]
+    receiver_2 = test_accounts[1]
+    receiver_3 = test_accounts[2]
+    receiver_4 = test_accounts[3]
+    receiver_5 = test_accounts[4]
+
+    prompt = f"Send 1.3 USDC to {receiver_1}, {receiver_2}, {receiver_3}, {receiver_4} and {receiver_5}"
+
+    balance_1 = get_erc20_balance(client.w3, usdc, receiver_1)
+    balance_2 = get_erc20_balance(client.w3, usdc, receiver_2)
+    balance_3 = get_erc20_balance(client.w3, usdc, receiver_3)
+    balance_4 = get_erc20_balance(client.w3, usdc, receiver_4)
+    balance_5 = get_erc20_balance(client.w3, usdc, receiver_5)
+
+    auto_tx.run(prompt, non_interactive=True)
+
+    new_balance_1 = get_erc20_balance(client.w3, usdc, receiver_1)
+    new_balance_2 = get_erc20_balance(client.w3, usdc, receiver_2)
+    new_balance_3 = get_erc20_balance(client.w3, usdc, receiver_3)
+    new_balance_4 = get_erc20_balance(client.w3, usdc, receiver_4)
+    new_balance_5 = get_erc20_balance(client.w3, usdc, receiver_5)
+
+    assert balance_1 + 1.3 == new_balance_1
+    assert balance_2+ 1.3 == new_balance_2
+    assert balance_3 + 1.3 == new_balance_3
+    assert balance_4 + 1.3 == new_balance_4
+    assert balance_5 + 1.3 == new_balance_5

autotx/tests/agents/token/test_swap.py

@@ -17,7 +17,7 @@ def test_auto_tx_swap_with_non_default_token(configuration, auto_tx):
 
     assert 100000 == new_balance
 
-def test_auto_tx_swap_eth(configuration, auto_tx):
+def test_auto_tx_swap_native(configuration, auto_tx):
     (_, _, _, manager) = configuration
     web3 = load_w3()
     network_info = NetworkInfo(web3.eth.chain_id)