Merge pull request #237 from polywrap/nerfzael/cost-and-prompt-improv…

…ements

Prompt and test improvements, cost tracking

.gitignore

@@ -56,3 +56,4 @@ cassetes
 
 /benchmarks
 /logs
+/costs

autotx/AutoTx.py

@@ -1,13 +1,15 @@
 from enum import Enum
+from datetime import datetime
 import json
+import os
 from textwrap import dedent
 from typing import Any, Dict, Optional, Callable
 from dataclasses import dataclass
 from autogen import Agent as AutogenAgent
 from termcolor import cprint
 from typing import Optional
 from autotx.autotx_agent import AutoTxAgent
-from autotx.helper_agents import clarifier, manager, user_proxy, verifier
+from autotx.helper_agents import clarifier, manager, user_proxy
 from autotx.utils.logging.Logger import Logger
 from autotx.utils.PreparedTx import PreparedTx
 from autotx.utils.ethereum import SafeManager
@@ -16,7 +18,15 @@
 @dataclass(kw_only=True)
 class Config:
     verbose: bool
-    logs_dir: Optional[str]
+    logs_dir: Optional[str] = None
+    log_costs: bool
+    max_rounds: int
+
+    def __init__(self, verbose: bool, logs_dir: Optional[str], max_rounds: Optional[int] = None, log_costs: Optional[bool] = None):
+        self.verbose = verbose
+        self.logs_dir = logs_dir
+        self.log_costs = log_costs if log_costs is not None else False
+        self.max_rounds = max_rounds if max_rounds is not None else 100
 
 @dataclass
 class PastRun:
@@ -33,6 +43,8 @@ class RunResult:
     chat_history_json: str
     transactions: list[PreparedTx]
     end_reason: EndReason
+    total_cost_without_cache: float
+    total_cost_with_cache: float
 
 class AutoTx:
     manager: SafeManager
@@ -41,6 +53,8 @@ class AutoTx:
     network: NetworkInfo
     get_llm_config: Callable[[], Optional[Dict[str, Any]]]
     agents: list[AutoTxAgent]
+    log_costs: bool
+    max_rounds: int
 
     def __init__(
         self,
@@ -58,11 +72,28 @@ def __init__(
             silent=not config.verbose
         )
         self.agents = agents
+        self.log_costs = config.log_costs
+        self.max_rounds = config.max_rounds
 
     def run(self, prompt: str, non_interactive: bool, summary_method: str = "last_msg") -> RunResult:
+        total_cost_without_cache: float = 0
+        total_cost_with_cache: float = 0
+
         while True:
             result = self.try_run(prompt, non_interactive, summary_method)
+            total_cost_without_cache += result.total_cost_without_cache
+            total_cost_with_cache += result.total_cost_with_cache
+
             if result.end_reason == EndReason.TERMINATE or non_interactive:
+                if self.log_costs:
+                    now = datetime.now()
+                    now_str = now.strftime('%Y-%m-%d-%H-%M-%S-') + str(now.microsecond)
+
+                    if not os.path.exists("costs"):
+                        os.makedirs("costs")
+                    with open(f"costs/{now_str}.txt", "w") as f:
+                        f.write(str(total_cost_without_cache))
+
                 return result
             else:
                 cprint("Prompt not supported. Please provide a new prompt.", "yellow")
@@ -98,19 +129,27 @@ def try_run(self, prompt: str, non_interactive: bool, summary_method: str = "las
             agents_information = self.get_agents_information(self.agents)
 
             user_proxy_agent = user_proxy.build(prompt, agents_information, self.get_llm_config)
-            clarifier_agent = clarifier.build(user_proxy_agent, agents_information, self.manager.address, self.network.chain_id.name, non_interactive, self.get_llm_config)
+            clarifier_agent = clarifier.build(user_proxy_agent, agents_information, non_interactive, self.get_llm_config)
 
             helper_agents: list[AutogenAgent] = [
                 user_proxy_agent,
-                verifier.build(self.get_llm_config),
                 clarifier_agent
             ]
 
             autogen_agents = [agent.build_autogen_agent(self, user_proxy_agent, self.get_llm_config()) for agent in self.agents]
 
-            manager_agent = manager.build(autogen_agents + helper_agents, self.get_llm_config)
-
-            chat = user_proxy_agent.initiate_chat(manager_agent, message=prompt, summary_method=summary_method)
+            manager_agent = manager.build(autogen_agents + helper_agents, self.max_rounds, self.get_llm_config)
+
+            chat = user_proxy_agent.initiate_chat(
+                manager_agent, 
+                message=dedent(
+                    f"""
+                    I am currently connected with the following wallet: {self.manager.address}, on network: {self.network.chain_id.name}
+                    My goal is: {prompt} 
+                    """
+                ), 
+                summary_method=summary_method
+            )
 
             if "ERROR:" in chat.summary:
                 error_message = chat.summary.replace("ERROR: ", "").replace("\n", "")
@@ -147,7 +186,7 @@ def try_run(self, prompt: str, non_interactive: bool, summary_method: str = "las
 
         chat_history = json.dumps(chat.chat_history, indent=4)
 
-        return RunResult(chat.summary, chat_history, transactions, EndReason.TERMINATE if is_goal_supported else EndReason.GOAL_NOT_SUPPORTED)
+        return RunResult(chat.summary, chat_history, transactions, EndReason.TERMINATE if is_goal_supported else EndReason.GOAL_NOT_SUPPORTED, float(chat.cost["usage_including_cached_inference"]["total_cost"]), float(chat.cost["usage_excluding_cached_inference"]["total_cost"]))
 
     def get_agents_information(self, agents: list[AutoTxAgent]) -> str:
         agent_descriptions = []

autotx/agents/ResearchTokensAgent.py

@@ -1,7 +1,7 @@
 
 import json
 from textwrap import dedent
-from typing import Annotated, Any, Callable, Optional, Type, Union, cast
+from typing import Annotated, Callable, Optional, Union, cast
 from web3 import Web3
 from autotx.AutoTx import AutoTx
 from gnosis.eth import EthereumNetworkNotSupported as ChainIdNotSupported
@@ -15,12 +15,19 @@
 name = "research-tokens"
 
 system_message = lambda autotx: dedent(f"""
-    You are an AI assistant. Assist the user (address: {autotx.manager.address}) in their task of researching tokens.
-    You are an expert in Ethereum tokens and can help users research tokens.
+    You are an AI assistant that's an expert in Ethereum tokens. Assist the user in their task of researching tokens.
     ONLY focus on the token research aspect of the user's goal and let other agents handle other tasks.
     You use the tools available to assist the user in their tasks.
     Retrieve token information, get token price, market cap, and price change percentage.
-    Always fetch available token categories before suggesting names, do not make them up.
+    BEFORE calling get_tokens_based_on_category always call get_available_categories to get the list of available categories.
+    You MUST keep in mind the network the user is on and if the request is for a specific network, all networks, or the current network (it could be implied).
+    If the user is interested in buying the tokens you're researching make sure you're searching them for his network.
+    If searching for tokens to buy, make sure to:
+    1. Get all the tokens you need (if fetching from multiple categories the tokens could overlap, so you might need to fetch more until you get the number you need)
+    2. Calculate how much of each token to buy
+    If a token does not have enough liquidity on the network, either:
+    - Search for another token (if it is within the user's goal)
+    - Or inform the user that the token is not available on the network (E.g. if the user's goal is to buy that specific token)
     """
 )
 

autotx/agents/SendTokensAgent.py

@@ -20,7 +20,7 @@
 name = "send-tokens"
 
 system_message = lambda autotx: dedent(f"""
-    You are an expert in Ethereum tokens (native and erc20) and can 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 (native and erc20) and can assist the user in their tasks by fetching balances and preparing transactions to send tokens.
     You are in a group of agents that will help the user achieve their goal.
     ONLY focus on the sending and balance aspect of the user's goal and let other agents handle other tasks.
     You use the tools available to assist the user in their tasks.