Merge pull request #164 from polywrap/chore/build-script

chore: build script

.github/workflows/ci.yaml

@@ -0,0 +1,28 @@
+name: ci
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v2
+
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v2
+      with:
+        python-version: "3.10"
+
+    - name: Install Poetry
+      run: curl -sSL https://install.python-poetry.org | python3 -
+
+    - name: Install dependencies
+      run: poetry install
+
+    - name: Check types
+      run: poetry run build-check

autotx/AutoTx.py

@@ -4,13 +4,14 @@
 from autogen import UserProxyAgent, AssistantAgent, GroupChat, GroupChatManager
 from termcolor import cprint
 from typing import Optional
-from autogen.io import IOStream
+from autogen.io import IOStream, IOConsole
 from autotx.autotx_agent import AutoTxAgent
 from autotx.utils.PreparedTx import PreparedTx
 from autotx.utils.agent.build_goal import build_goal
 from autotx.utils.ethereum import SafeManager
 from autotx.utils.ethereum.networks import NetworkInfo
-from autotx.utils.io_silent import IOConsole, IOSilent
+from autotx.utils.io_silent import IOSilent
+
 
 @dataclass(kw_only=True)
 class Config:
@@ -22,7 +23,6 @@ class AutoTx:
     transactions: list[PreparedTx] = []
     network: NetworkInfo
     get_llm_config: Callable[[], Optional[Dict[str, Any]]]
-    user_proxy: UserProxyAgent
     agents: list[AutoTxAgent]
 
     def __init__(
@@ -36,7 +36,7 @@ def __init__(
             self.config = config
         self.agents = agents
 
-    def run(self, prompt: str, non_interactive: bool, silent: bool = False):
+    def run(self, prompt: str, non_interactive: bool, silent: bool = False) -> None:
         print("Running AutoTx with the following prompt: ", prompt)
 
         user_proxy = UserProxyAgent(

autotx/agents/ExampleAgent.py

@@ -18,7 +18,7 @@ class ExampleTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[float, str], str]:
         def run(
             amount: Annotated[float, "Amount of something."],
             receiver: Annotated[str, "The receiver of something."]

autotx/agents/ResearchTokensAgent.py

@@ -1,7 +1,7 @@
 
 import json
 from textwrap import dedent
-from typing import Annotated, Callable, List, Optional, Union
+from typing import Annotated, Any, Callable, Optional, Type, Union, cast
 from web3 import Web3
 from autotx.AutoTx import AutoTx
 from gnosis.eth import EthereumNetworkNotSupported as ChainIdNotSupported
@@ -31,14 +31,14 @@
     ChainId.GNOSIS: "xdai",
 }
 
-def get_coingecko():
+def get_coingecko() -> CoinGeckoDemoClient:
     return CoinGeckoDemoClient(api_key=COINGECKO_API_KEY)
 
 
 def get_tokens_and_filter_per_network(
     network_name: str,
-) -> dict[str, Union[str, dict[str, str]]]:
-    network = ChainId[network_name]
+) -> list[dict[str, Union[str, dict[str, str]]]]:
+    network = ChainId[network_name] # type: ignore
     coingecko_network_key = COINGECKO_NETWORKS_TO_SUPPORTED_NETWORKS_MAP.get(network)
     if coingecko_network_key == None:
         raise ChainIdNotSupported(f"Network {network_name} not supported")
@@ -50,11 +50,12 @@ def get_tokens_and_filter_per_network(
         if coingecko_network_key in token["platforms"]
     ]
 
-def filter_token_list_by_network(tokens: list[dict[str, str]], network_name: str):
+def filter_token_list_by_network(tokens: list[dict[str, str]], network_name: str) -> list[dict[str, Union[str, dict[str, str]]]]:
     tokens_in_category_map = {category["id"]: category for category in tokens}
     filtered_tokens_map = {
         token["id"]: token for token in get_tokens_and_filter_per_network(network_name)
     }
+
     return [
         {
             **tokens_in_category_map[token["id"]],
@@ -65,16 +66,17 @@ def filter_token_list_by_network(tokens: list[dict[str, str]], network_name: str
     ]
 
 def add_tokens_address_if_not_in_registry(
-    tokens_in_category: list[dict[str, str]],
-    tokens: list[str, str],
+    tokens_in_category: list[dict[str, Union[str, dict[str, str]]]],
+    tokens: dict[str, str],
     current_network: str,
-):
+) -> None:
     for token_with_address in tokens_in_category:
-        token_symbol = token_with_address["symbol"].lower()
+        token_symbol = cast(str, token_with_address["symbol"]).lower()
         token_not_added = token_symbol not in tokens
         if token_not_added:
+            platforms = cast(dict[str, str], token_with_address["platforms"]) 
             tokens[token_symbol] = Web3.to_checksum_address(
-                token_with_address["platforms"][current_network]
+                platforms[current_network]
             )
 
 class GetTokenInformationTool(AutoTxTool):
@@ -85,7 +87,7 @@ class GetTokenInformationTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str], str]:
         def run(
             token_id: Annotated[str, "ID of token"]
         ) -> str:
@@ -140,7 +142,7 @@ class SearchTokenTool(AutoTxTool):
     name: str = "search_token"
     description: str = "Search token based on its symbol. It will return the ID of tokens with the largest market cap"
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str, bool], str]:
         def run(
             token_symbol: Annotated[str, "Symbol of token to search"],
             retrieve_duplicate: Annotated[bool, "Set to true to retrieve all instances of tokens sharing the same symbol, indicating potential duplicates. By default, it is False, meaning only a single, most relevant token is retrieved unless duplication is explicitly requested."]
@@ -161,7 +163,7 @@ class GetAvailableCategoriesTool(AutoTxTool):
     name: str = "get_available_categories"
     description: str = "Retrieve all available category ids of tokens"
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[], str]:
         def run() -> str:
             print("Fetching available token categories")
 
@@ -174,7 +176,7 @@ class GetTokensBasedOnCategoryTool(AutoTxTool):
     name: str = "get_tokens_based_on_category"
     description: str = "Retrieve all tokens with their respective information (symbol, market cap, price change percentages and total traded volume in the last 24 hours) from a given category"
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str, str, int, str, Optional[str]], str]:
         def run(
             category: Annotated[str, "Category to retrieve tokens"],
             sort_by: Annotated[str, "Sort tokens by field. It can be: 'volume_desc' | 'volume_asc' | 'market_cap_desc' | 'market_cap_asc'. 'market_cap_desc' is the default"],
@@ -207,11 +209,11 @@ def run(
                     autotx.network.chain_id
                 )
                 asked_network = COINGECKO_NETWORKS_TO_SUPPORTED_NETWORKS_MAP.get(
-                    ChainId[network_name]
+                    ChainId[network_name] # type: ignore
                 )
                 if current_network == asked_network:
                     add_tokens_address_if_not_in_registry(
-                        tokens_in_category, autotx.network.tokens, current_network
+                        tokens_in_category, autotx.network.tokens, cast(str, current_network)
                     )
 
             interval = (
@@ -238,10 +240,10 @@ class GetExchangesWhereTokenCanBeTradedTool(AutoTxTool):
     name: str = "get_exchanges_where_token_can_be_traded"
     description: str = "Retrieve exchanges where token can be traded"
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str], list[str]]:
         def run(
             token_id: Annotated[str, "ID of token"]
-        ) -> List[str]:
+        ) -> list[str]:
             print(f"Fetching exchanges where token ({token_id}) can be traded")
 
             tickers = get_coingecko().coins.get_tickers(id=token_id)["tickers"]

autotx/agents/SendTokensAgent.py

@@ -1,5 +1,7 @@
 from textwrap import dedent
 from typing import Annotated, Callable
+
+from web3 import Web3
 from autotx.AutoTx import AutoTx
 from autotx.autotx_agent import AutoTxAgent
 from autotx.autotx_tool import AutoTxTool
@@ -32,7 +34,7 @@ class TransferETHTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    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"]
@@ -62,13 +64,13 @@ class TransferERC20Tool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[float, str, 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"],
             token: Annotated[str, "Symbol of token to transfer"]
         ) -> str:
-            token_address = autotx.network.tokens[token.lower()]
+            token_address = Web3.to_checksum_address(autotx.network.tokens[token.lower()])
             web3 = load_w3()
 
             receiver_addr = ETHAddress(receiver, web3)
@@ -93,7 +95,7 @@ class GetETHBalanceTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str], float]:
         def run(
             owner: Annotated[str, "The owner's address or ENS domain"]
         ) -> float:
@@ -117,7 +119,7 @@ class GetERC20BalanceTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str, str], float]:
         def run(
             token: Annotated[str, "Token symbol of erc20"],
             owner: Annotated[str, "The token owner's address or ENS domain"]

autotx/agents/SwapTokensAgent.py

@@ -1,3 +1,4 @@
+from decimal import Decimal
 from textwrap import dedent
 from typing import Annotated, Callable
 from autotx.AutoTx import AutoTx
@@ -34,7 +35,7 @@
     """
 )
 
-def get_tokens_address(token_in: str, token_out: str, network_info: NetworkInfo):
+def get_tokens_address(token_in: str, token_out: str, network_info: NetworkInfo) -> tuple[str, str]:
     token_in = token_in.lower()
     token_out = token_out.lower()
 
@@ -59,7 +60,7 @@ class SwapTool(AutoTxTool):
         """
     )
 
-    def build_tool(self, autotx: AutoTx) -> Callable:
+    def build_tool(self, autotx: AutoTx) -> Callable[[str, str], str]:
         def run(
             token_to_sell: Annotated[str, "Token to sell. E.g. '10 USDC' or just 'USDC'"],
             token_to_buy: Annotated[str, "Token to buy. E.g. '10 USDC' or just 'USDC'"],
@@ -89,7 +90,7 @@ def run(
 
             swap_transactions = build_swap_transaction(
                 autotx.manager.client,
-                float(exact_amount),
+                Decimal(exact_amount),
                 token_in_address,
                 token_out_address,
                 autotx.manager.address.hex,

autotx/autotx_agent.py

@@ -10,7 +10,7 @@ class AutoTxAgent():
     tools: list['AutoTxTool']
     tool_descriptions: list[str]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.tool_descriptions = [
             f"{tool.name}: {tool.description}" for tool in self.tools
         ]
@@ -20,7 +20,7 @@ def build_autogen_agent(self, autotx: 'AutoTx', user_proxy: autogen.UserProxyAge
         if isinstance(self.system_message, str):
             system_message = self.system_message
         else:
-            get_system_message = self.system_message.__func__
+            get_system_message = self.system_message.__func__ # type: ignore
             system_message = get_system_message(autotx)
 
         agent = autogen.AssistantAgent(

autotx/autotx_tool.py

@@ -1,4 +1,4 @@
-from typing import Callable, TYPE_CHECKING
+from typing import Any, Callable, TYPE_CHECKING
 from autogen import Agent, UserProxyAgent
 if TYPE_CHECKING:
   from autotx.AutoTx import AutoTx
@@ -7,12 +7,12 @@ class AutoTxTool:
   name: str
   description: str
 
-  def register_tool(self, autotx: 'AutoTx', agent: Agent, user_proxy: UserProxyAgent):
+  def register_tool(self, autotx: 'AutoTx', agent: Agent, user_proxy: UserProxyAgent) -> None:
     tool = self.build_tool(autotx)
     # Register the tool signature with the assistant agent.
     agent.register_for_llm(name=self.name, description=self.description)(tool)
     # Register the tool function with the user proxy agent.
     user_proxy.register_for_execution(name=self.name)(tool)
 
-  def build_tool(self, autotx: 'AutoTx') -> Callable:
+  def build_tool(self, autotx: 'AutoTx') -> Callable[..., Any]:
     raise NotImplementedError

autotx/build_check.py

@@ -0,0 +1,10 @@
+import subprocess
+import sys
+
+def run() -> None:
+    # Run mypy check
+    result = subprocess.run(["mypy", "."], capture_output=True)
+    print(result.stdout.decode())
+    if result.returncode != 0:
+        print("Type checking failed")
+        sys.exit(1)

autotx/chain_fork.py

@@ -4,7 +4,7 @@
 
 container_name = "autotx_chain_fork"
 
-def start():
+def start() -> None:
     delete_cached_safe_address()
 
     build = subprocess.run(
@@ -32,5 +32,5 @@ def start():
         check=True,
     )
 
-def stop():
+def stop() -> None:
     subprocess.run(["docker", "container", "rm", container_name, "-f"], check=True)