Release 0.1.4

autotx/AutoTx.py

@@ -4,14 +4,14 @@
 import json
 import os
 from textwrap import dedent
-from typing import Any, Dict, Optional, Callable
-from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Callable, Union
+from dataclasses import dataclass
 from autogen import Agent as AutogenAgent, ModelClient
 from termcolor import cprint
 from typing import Optional
+from autogen import AssistantAgent
 
 from web3 import Web3
-from autotx import models
 from autotx.autotx_agent import AutoTxAgent
 from autotx.helper_agents import clarifier, manager, user_proxy
 from autotx.intents import Intent
@@ -34,14 +34,25 @@ class Config:
     max_rounds: int
     get_llm_config: Callable[[], Optional[Dict[str, Any]]]
     custom_model: Optional[CustomModel] = None
-
-    def __init__(self, verbose: bool, get_llm_config: Callable[[], Optional[Dict[str, Any]]], logs_dir: Optional[str], max_rounds: Optional[int] = None, log_costs: Optional[bool] = None, custom_model: Optional[CustomModel] = None):
+    on_agent_message: Optional[Callable[[str, str, Any], None]] = None
+
+    def __init__(self, 
+        verbose: bool, 
+        get_llm_config: Callable[[], 
+        Optional[Dict[str, Any]]], 
+        logs_dir: Optional[str], 
+        max_rounds: Optional[int] = None, 
+        log_costs: Optional[bool] = None, 
+        custom_model: Optional[CustomModel] = None,
+        on_agent_message: Optional[Callable[[str, str, Any], None]] = None
+    ):
         self.verbose = verbose
         self.get_llm_config = get_llm_config
         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
+        self.max_rounds = max_rounds if max_rounds is not None else 50
         self.custom_model = custom_model
+        self.on_agent_message = on_agent_message
 
 @dataclass
 class PastRun:
@@ -78,6 +89,7 @@ class AutoTx:
     info_messages: list[str]
     verbose: bool
     on_notify_user: Callable[[str], None] | None
+    on_agent_message: Optional[Callable[[AssistantAgent, Union[Dict[str, Any], str], AutogenAgent, bool], Union[Dict[str, Any], str]]]
 
     def __init__(
         self,
@@ -109,6 +121,7 @@ def __init__(
         self.info_messages = []
         self.on_notify_user = on_notify_user
         self.custom_model = config.custom_model
+        self.on_agent_message = build_on_message_hook(config.on_agent_message) if config.on_agent_message else None
 
     def run(self, prompt: str, non_interactive: bool, summary_method: str = "last_msg") -> RunResult:
         return asyncio.run(self.a_run(prompt, non_interactive, summary_method))    
@@ -194,7 +207,7 @@ async def try_run(self, prompt: str, non_interactive: bool, summary_method: str
 
             agents_information = self.get_agents_information(self.agents)
 
-            user_proxy_agent = user_proxy.build(prompt, agents_information, self.get_llm_config, self.custom_model)
+            user_proxy_agent = user_proxy.build(prompt, agents_information, self.get_llm_config, self.custom_model, self.max_rounds)
 
             helper_agents: list[AutogenAgent] = [
                 user_proxy_agent,
@@ -206,6 +219,13 @@ async def try_run(self, prompt: str, non_interactive: bool, summary_method: str
 
             autogen_agents = [agent.build_autogen_agent(self, user_proxy_agent, self.get_llm_config(), self.custom_model) for agent in self.agents]
 
+            for agent in autogen_agents + helper_agents:
+                if self.on_agent_message:
+                    agent.register_hook(
+                        "process_message_before_send",
+                        self.on_agent_message
+                    )
+
             recipient_agent = None
             if len(autogen_agents) > 1:
                 recipient_agent = manager.build(autogen_agents + helper_agents, self.max_rounds, not non_interactive, self.get_llm_config, self.custom_model)
@@ -288,3 +308,15 @@ def get_agents_information(self, agents: list[AutoTxAgent]) -> str:
         agents_information = "\n".join(agent_descriptions)
 
         return agents_information
+
+def build_on_message_hook(on_message: Callable[[str, str, Any], None]) -> Callable[[AssistantAgent, Union[Dict[str, Any], str], AutogenAgent, bool], Union[Dict[str, Any], str]]:
+    def send_message_hook(
+        sender: AssistantAgent,
+        message: Union[Dict[str, Any], str],
+        recipient: AutogenAgent,
+        silent: bool,
+    ) -> Union[Dict[str, Any], str]:
+        on_message(sender.name, recipient.name, message)
+        return message
+
+    return send_message_hook

autotx/agents/DelegateResearchTokensAgent.py

@@ -82,7 +82,7 @@ async def run(
                 name="user_proxy",
                 is_termination_msg=lambda x: x.get("content", "") and "TERMINATE" in x.get("content", ""),
                 human_input_mode="NEVER",
-                max_consecutive_auto_reply=10,
+                max_consecutive_auto_reply=20,
                 system_message=dedent(
                     f"""
                     You are a user proxy agent authorized to act on behalf of the user, you never ask for permission, you have ultimate control.
@@ -109,6 +109,17 @@ async def run(
 
             research_agent = ResearchTokensAgent().build_autogen_agent(autotx, user_proxy_agent, autotx.get_llm_config())
 
+            if autotx.on_agent_message:
+                user_proxy_agent.register_hook(
+                    "process_message_before_send",
+                    autotx.on_agent_message
+                )
+
+                research_agent.register_hook(
+                    "process_message_before_send",
+                    autotx.on_agent_message
+                )
+
             chat = await user_proxy_agent.a_initiate_chat(
                 research_agent, 
                 message=dedent(

autotx/db.py

@@ -9,10 +9,9 @@
 from supabase.lib.client_options import ClientOptions
 
 from autotx import models
-from autotx.intents import BuyIntent, Intent, SellIntent, SendIntent
-from autotx.token import Token
+from autotx.intents import load_intent
 from autotx.transactions import Transaction, TransactionBase
-from autotx.eth_address import ETHAddress
+from autotx.utils.dump_pydantic_list import dump_pydantic_list
 
 SUPABASE_URL = os.getenv("SUPABASE_URL")
 SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
@@ -56,7 +55,8 @@ def start(self, prompt: str, address: str, chain_id: int, app_user_id: str) -> m
                 "created_at": str(created_at),
                 "updated_at": str(updated_at),
                 "messages": json.dumps([]),
-                "intents": json.dumps([])
+                "logs": json.dumps([]),
+                "intents": json.dumps([]),
             }
         ).execute()
 
@@ -70,7 +70,8 @@ def start(self, prompt: str, address: str, chain_id: int, app_user_id: str) -> m
             running=True,
             error=None,
             messages=[],
-            intents=[]
+            logs=[],
+            intents=[],
         )
 
     def stop(self, task_id: str) -> None:
@@ -86,16 +87,15 @@ def stop(self, task_id: str) -> None:
     def update(self, task: models.Task) -> None:
         client = get_db_client("public")
 
-        intents = [json.loads(intent.json()) for intent in task.intents]
-
         client.table("tasks").update(
             {
                 "prompt": task.prompt,
                 "running": task.running,
                 "updated_at": str(datetime.utcnow()),
                 "messages": json.dumps(task.messages),
                 "error": task.error,
-                "intents": json.dumps(intents)
+                "logs": dump_pydantic_list(task.logs if task.logs else []),
+                "intents": dump_pydantic_list(task.intents)
             }
         ).eq("id", task.id).eq("app_id", self.app_id).execute()
 
@@ -113,43 +113,6 @@ def get(self, task_id: str) -> models.Task | None:
 
         task_data = result.data[0]
 
-        def load_intent(intent_data: dict[str, Any]) -> Intent:
-            if intent_data["type"] == "send":
-                return SendIntent.create(
-                    receiver=ETHAddress(intent_data["receiver"]),
-                    token=Token(
-                        symbol=intent_data["token"]["symbol"],
-                        address=intent_data["token"]["address"]
-                    ),
-                    amount=intent_data["amount"]
-                )
-            elif intent_data["type"] == "buy":
-                return BuyIntent.create(
-                    from_token=Token(
-                        symbol=intent_data["from_token"]["symbol"],
-                        address=intent_data["from_token"]["address"]
-                    ),
-                    to_token=Token(
-                        symbol=intent_data["to_token"]["symbol"],
-                        address=intent_data["to_token"]["address"]
-                    ),
-                    amount=intent_data["amount"]
-                )
-            elif intent_data["type"] == "sell":
-                return SellIntent.create(
-                    from_token=Token(
-                        symbol=intent_data["from_token"]["symbol"],
-                        address=intent_data["from_token"]["address"]
-                    ),
-                    to_token=Token(
-                        symbol=intent_data["to_token"]["symbol"],
-                        address=intent_data["to_token"]["address"]
-                    ),
-                    amount=intent_data["amount"]
-                )
-            else:
-                raise Exception(f"Unknown intent type: {intent_data['type']}")
-
         return models.Task(
             id=task_data["id"],
             prompt=task_data["prompt"],
@@ -160,6 +123,7 @@ def load_intent(intent_data: dict[str, Any]) -> Intent:
             running=task_data["running"],
             error=task_data["error"],
             messages=json.loads(task_data["messages"]),
+            logs=[models.TaskLog(**log) for log in json.loads(task_data["logs"])] if task_data["logs"] else None,
             intents=[load_intent(intent) for intent in json.loads(task_data["intents"])]
         )
 
@@ -182,7 +146,8 @@ def get_all(self) -> list[models.Task]:
                     running=task_data["running"],
                     error=task_data["error"],
                     messages=json.loads(task_data["messages"]),
-                    intents=json.loads(task_data["intents"])
+                    logs=[models.TaskLog(**log) for log in json.loads(task_data["logs"])] if task_data["logs"] else None,
+                    intents=[load_intent(intent) for intent in json.loads(task_data["intents"])]
                 )
             )
 
@@ -318,6 +283,21 @@ def submit_transactions(app_id: str, app_user_id: str, submitted_batch_id: str)
         .eq("app_user_id", app_user_id) \
         .eq("id", submitted_batch_id) \
         .execute()
+
+def add_task_error(context: str, app_id: str, app_user_id: str, task_id: str, message: str) -> None:
+    client = get_db_client("public")
+
+    created_at = datetime.utcnow()
+    client.table("task_errors").insert(
+        {
+            "context": context,
+            "app_id": app_id,
+            "task_id": task_id,
+            "app_user_id": app_user_id,
+            "message": message,
+            "created_at": str(created_at)
+        }
+    ).execute()
 
 class SubmittedBatch(BaseModel):
     id: str
@@ -358,6 +338,21 @@ def get_submitted_batches(app_id: str, task_id: str) -> list[SubmittedBatch]:
 
     return batches
 
+def get_task_logs(task_id: str) -> list[models.TaskLog] | None:
+    client = get_db_client("public")
+
+    result = client.table("tasks") \
+        .select("logs") \
+        .eq("id", task_id) \
+        .execute()
+
+    if len(result.data) == 0:
+        return None
+
+    task_data = result.data[0]
+
+    return [models.TaskLog(**log) for log in json.loads(task_data["logs"])] if task_data["logs"] else []
+
 def create_app(name: str, api_key: str) -> models.App:
     client = get_db_client("public")
 

autotx/helper_agents/user_proxy.py

@@ -5,12 +5,12 @@
 if TYPE_CHECKING:
     from autotx.AutoTx import CustomModel
 
-def build(user_prompt: str, agents_information: str, get_llm_config: Callable[[], Optional[Dict[str, Any]]], custom_model: Optional['CustomModel']) -> UserProxyAgent:
+def build(user_prompt: str, agents_information: str, get_llm_config: Callable[[], Optional[Dict[str, Any]]], custom_model: Optional['CustomModel'], max_rounds: int) -> UserProxyAgent:
     user_proxy = UserProxyAgent(
         name="user_proxy",
         is_termination_msg=lambda x: x.get("content", "") and "TERMINATE" in x.get("content", ""),
         human_input_mode="NEVER",
-        max_consecutive_auto_reply=4 if custom_model else 10,
+        max_consecutive_auto_reply=4 if custom_model else max_rounds,
         system_message=dedent(
             f"""
             You are a user proxy agent authorized to act on behalf of the user, you never ask for permission, you have ultimate control.

autotx/intents.py

@@ -122,3 +122,40 @@ async def build_transactions(self, web3: Web3, network: NetworkInfo, smart_walle
         return transactions
 
 Intent = Union[SendIntent, BuyIntent, SellIntent]
+
+def load_intent(intent_data: dict[str, Any]) -> Intent:
+    if intent_data["type"] == "send":
+        return SendIntent.create(
+            receiver=ETHAddress(intent_data["receiver"]),
+            token=Token(
+                symbol=intent_data["token"]["symbol"],
+                address=intent_data["token"]["address"]
+            ),
+            amount=intent_data["amount"]
+        )
+    elif intent_data["type"] == "buy":
+        return BuyIntent.create(
+            from_token=Token(
+                symbol=intent_data["from_token"]["symbol"],
+                address=intent_data["from_token"]["address"]
+            ),
+            to_token=Token(
+                symbol=intent_data["to_token"]["symbol"],
+                address=intent_data["to_token"]["address"]
+            ),
+            amount=intent_data["amount"]
+        )
+    elif intent_data["type"] == "sell":
+        return SellIntent.create(
+            from_token=Token(
+                symbol=intent_data["from_token"]["symbol"],
+                address=intent_data["from_token"]["address"]
+            ),
+            to_token=Token(
+                symbol=intent_data["to_token"]["symbol"],
+                address=intent_data["to_token"]["address"]
+            ),
+            amount=intent_data["amount"]
+        )
+    else:
+        raise Exception(f"Unknown intent type: {intent_data['type']}")

autotx/models.py

@@ -1,9 +1,14 @@
 from pydantic import BaseModel
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 from datetime import datetime
 
 from autotx.intents import Intent
 
+class TaskLog(BaseModel):
+    type: str
+    obj: str
+    created_at: datetime
+
 class Task(BaseModel):
     id: str
     prompt: str
@@ -14,8 +19,17 @@ class Task(BaseModel):
     error: str | None
     running: bool
     messages: List[str]
+    logs: List[TaskLog] | None
     intents: List[Intent]
 
+class TaskError(BaseModel):
+    id: str
+    message: str
+    task_id: str
+    app_id: str
+    app_user_id: str
+    created_at: datetime
+
 class App(BaseModel):
     id: str
     name: str