🐛 Fix syntax and parsing errors

tools/dao.py

@@ -1,140 +1,212 @@
 from typing import Optional, Type, List
+from pydantic import BaseModel, Field
 from crewai_tools import BaseTool
 from .bun import BunScriptRunner
-from pydantic import BaseModel, Field
 
+# Schema definitions
+class ExecutorListToolSchema(BaseModel):
+    """Input schema for listing executor contracts."""
+    pass
 
 class ExecutorDeployToolSchema(BaseModel):
-    """Input schema for ExecutorDeployTool."""
-
+    """Input schema for deploying executor contracts."""
     name: str = Field(..., description="Name of the DAO")
     extensions: Optional[List[str]] = Field(
-        default=[], description="List of extension contracts to include"
+        default=[], description="Extension contracts to include"
     )
     include_deployer: bool = Field(
         default=False, description="Whether to include deployer in extensions"
     )
+    fee: Optional[int] = Field(
+        default=400000, description="Transaction fee in microSTX"
+    )
+
+class ExecutorSetExtensionToolSchema(BaseModel):
+    """Input schema for setting extension status."""
+    extension: str = Field(..., description="Extension contract ID")
+    status: str = Field(..., description="'enable' or 'disable'")
+    executor: str = Field(..., description="Executor contract ID")
+    fee: Optional[int] = Field(
+        default=10000, description="Transaction fee in microSTX"
+    )
 
+class TreasuryListToolSchema(BaseModel):
+    """Input schema for listing treasury contracts."""
+    pass
 
-class ExecutorDeployTool(BaseTool):
-    name: str = "Deploy a new DAO executor contract"
-    description: str = "Deploy a new executor contract that manages the DAO"
-    args_schema: Type[BaseModel] = ExecutorDeployToolSchema
+class TreasuryDeployToolSchema(BaseModel):
+    """Input schema for deploying treasury contracts."""
+    name: str = Field(..., description="Name of the Treasury")
+    dao_contract_id: Optional[str] = Field(
+        None, description="Associated DAO contract ID"
+    )
+    fee: Optional[int] = Field(
+        default=400000, description="Transaction fee in microSTX"
+    )
+
+class TreasuryDepositToolSchema(BaseModel):
+    """Input schema for depositing STX."""
+    treasury_id: str = Field(..., description="Treasury contract ID")
+    amount: int = Field(..., description="Amount in microSTX")
+    fee: Optional[int] = Field(
+        default=10000, description="Transaction fee in microSTX"
+    )
+
+class TreasuryWithdrawToolSchema(BaseModel):
+    """Input schema for withdrawing STX."""
+    treasury_id: str = Field(..., description="Treasury contract ID")
+    amount: int = Field(..., description="Amount in microSTX")
+    recipient: Optional[str] = Field(
+        None, description="Optional recipient address"
+    )
+    fee: Optional[int] = Field(
+        default=10000, description="Transaction fee in microSTX"
+    )
+
+# Base Tool
+class DAOBaseTool(BaseTool):
     account_index: Optional[str] = None
 
-    def __init__(self, account_index: str, **kwargs):
+    def __init__(self, account_index: str = "0", **kwargs):
         super().__init__(**kwargs)
         self.account_index = account_index
 
-    def _run(
-        self,
-        name: str,
-        extensions: List[str] = [],
-        include_deployer: bool = False,
-    ) -> str:
-        extensions_str = " ".join([f"-e {ext}" for ext in extensions])
-        deployer_flag = "-d" if include_deployer else ""
+# Executor Tools
+class ExecutorListTool(DAOBaseTool):
+    name: str = "DAO: List Executors"
+    description: str = "List all executor contracts"
+    args_schema: Type[BaseModel] = ExecutorListToolSchema
+
+    def _run(self) -> str:
         return BunScriptRunner.bun_run(
             self.account_index,
             "stacks-dao",
             "cli.ts",
             "executor",
-            "deploy",
-            "-n", name,
-            extensions_str,
-            deployer_flag
+            "list"
         )
 
+class ExecutorDeployTool(DAOBaseTool):
+    name: str = "DAO: Deploy Executor"
+    description: str = "Deploy a new executor contract"
+    args_schema: Type[BaseModel] = ExecutorDeployToolSchema
 
-class TreasuryDeployToolSchema(BaseModel):
-    """Input schema for TreasuryDeployTool."""
-
-    name: str = Field(..., description="Name of the DAO")
-    dao_contract_id: str = Field(..., description="Contract ID of the executor DAO")
-
+    def _run(self, name: str, extensions: List[str] = [], 
+             include_deployer: bool = False, fee: int = 400000) -> str:
+        args = ["executor", "deploy", "-n", name]
+        if extensions:
+            args.extend(["-e", *extensions])
+        if include_deployer:
+            args.append("-d")
+        if fee != 400000:
+            args.extend(["-f", str(fee)])
+
+        return BunScriptRunner.bun_run(
+            self.account_index,
+            "stacks-dao",
+            "cli.ts",
+            *args
+        )
 
-class TreasuryDeployTool(BaseTool):
-    name: str = "Deploy a new DAO treasury contract"
-    description: str = "Deploy a new treasury contract for managing DAO assets"
-    args_schema: Type[BaseModel] = TreasuryDeployToolSchema
-    account_index: Optional[str] = None
+class ExecutorSetExtensionTool(DAOBaseTool):
+    name: str = "DAO: Set Extension Status"
+    description: str = "Enable or disable an extension contract"
+    args_schema: Type[BaseModel] = ExecutorSetExtensionToolSchema
+
+    def _run(self, extension: str, status: str, executor: str, 
+             fee: int = 10000) -> str:
+        args = [
+            "executor", "set-extension",
+            "-e", extension,
+            "-s", status,
+            "-x", executor
+        ]
+        if fee != 10000:
+            args.extend(["-f", str(fee)])
+
+        return BunScriptRunner.bun_run(
+            self.account_index,
+            "stacks-dao",
+            "cli.ts",
+            *args
+        )
 
-    def __init__(self, account_index: str, **kwargs):
-        super().__init__(**kwargs)
-        self.account_index = account_index
+# Treasury Tools
+class TreasuryListTool(DAOBaseTool):
+    name: str = "DAO: List Treasuries"
+    description: str = "List all treasury contracts"
+    args_schema: Type[BaseModel] = TreasuryListToolSchema
 
-    def _run(
-        self,
-        name: str,
-        dao_contract_id: str,
-    ) -> str:
+    def _run(self) -> str:
         return BunScriptRunner.bun_run(
             self.account_index,
             "stacks-dao",
             "cli.ts",
             "treasury",
-            "deploy",
-            "-n", name,
-            "-d", dao_contract_id
+            "list"
         )
 
+class TreasuryDeployTool(DAOBaseTool):
+    name: str = "DAO: Deploy Treasury"
+    description: str = "Deploy a new treasury contract"
+    args_schema: Type[BaseModel] = TreasuryDeployToolSchema
 
-class TreasuryDepositToolSchema(BaseModel):
-    """Input schema for TreasuryDepositTool."""
-
-    treasury_id: str = Field(..., description="Contract ID of the treasury")
-    amount: int = Field(..., description="Amount in microSTX to deposit")
-
+    def _run(self, name: str, dao_contract_id: Optional[str] = None, 
+             fee: int = 400000) -> str:
+        args = ["treasury", "deploy", "-n", name]
+        if dao_contract_id:
+            args.extend(["-d", dao_contract_id])
+        if fee != 400000:
+            args.extend(["-f", str(fee)])
+
+        return BunScriptRunner.bun_run(
+            self.account_index,
+            "stacks-dao",
+            "cli.ts",
+            *args
+        )
 
-class TreasuryDepositTool(BaseTool):
-    name: str = "Deposit STX into DAO treasury"
-    description: str = "Deposit STX tokens into the DAO treasury"
+class TreasuryDepositTool(DAOBaseTool):
+    name: str = "DAO: Deposit STX"
+    description: str = "Deposit STX into treasury"
     args_schema: Type[BaseModel] = TreasuryDepositToolSchema
-    account_index: Optional[str] = None
-
-    def __init__(self, account_index: str, **kwargs):
-        super().__init__(**kwargs)
-        self.account_index = account_index
 
-    def _run(
-        self,
-        treasury_id: str,
-        amount: int,
-    ) -> str:
+    def _run(self, treasury_id: str, amount: int, fee: int = 10000) -> str:
+        args = [
+            "treasury", "deposit-stx",
+            "-t", treasury_id,
+            "-a", str(amount)
+        ]
+        if fee != 10000:
+            args.extend(["-f", str(fee)])
+
         return BunScriptRunner.bun_run(
             self.account_index,
             "stacks-dao",
             "cli.ts",
-            "treasury",
-            "deposit-stx",
-            "-t", treasury_id,
-            "-a", str(amount),
+            *args
         )
 
-
-class TreasuryWithdrawTool(BaseTool):
-    name: str = "Withdraw STX from DAO treasury"
-    description: str = "Withdraw STX tokens from the DAO treasury"
+class TreasuryWithdrawTool(DAOBaseTool):
+    name: str = "DAO: Withdraw STX"
+    description: str = "Withdraw STX from treasury"
     args_schema: Type[BaseModel] = TreasuryWithdrawToolSchema
-    account_index: Optional[str] = None
 
-    def __init__(self, account_index: str, **kwargs):
-        super().__init__(**kwargs)
-        self.account_index = account_index
-
-    def _run(
-        self,
-        treasury_id: str,
-        amount: int,
-        recipient: str,
-    ) -> str:
+    def _run(self, treasury_id: str, amount: int, 
+             recipient: Optional[str] = None, fee: int = 10000) -> str:
+        args = [
+            "treasury", "withdraw-stx",
+            "-t", treasury_id,
+            "-a", str(amount)
+        ]
+        if recipient:
+            args.extend(["-r", recipient])
+        if fee != 10000:
+            args.extend(["-f", str(fee)])
+
         return BunScriptRunner.bun_run(
             self.account_index,
             "stacks-dao",
             "cli.ts",
-            "treasury",
-            "withdraw-stx",
-            "-t", treasury_id,
-            "-a", str(amount)
-            "-r", recipient
+            *args
         )

tools/tools_factory.py

@@ -17,7 +17,10 @@
 )
 from .contracts import ContractSIP10DeployTool, ContractSIP10SendTool, ContractSIP10InfoTool
 from .dao import (
+    ExecutorListTool,
     ExecutorDeployTool,
+    ExecutorSetExtensionTool,
+    TreasuryListTool,
     TreasuryDeployTool,
     TreasuryDepositTool,
     TreasuryWithdrawTool
@@ -27,8 +30,13 @@
 def initialize_tools(account_index: str = "0"):
     """
     Initialize and return a dictionary of available tools.
+
+    Args:
+        account_index (str): The account index to use for tools that require it.
+
+    Returns:
+        dict: A dictionary mapping tool names to tool instances.
     """
-    # this will be exposed by an endpoint for the frontend to get the available tools
     return {
         # Existing tools
         "alex_get_price_history": AlexGetPriceHistory(),
@@ -54,8 +62,13 @@ def initialize_tools(account_index: str = "0"):
         "fetch_contract_code": FetchContractCodeTool(),
         "get_btc_data": GetBitcoinData(),
 
-        # DAO tools
+        # DAO Executor tools
+        "dao_executor_list": ExecutorListTool(account_index),
         "dao_executor_deploy": ExecutorDeployTool(account_index),
+        "dao_executor_set_extension": ExecutorSetExtensionTool(account_index),
+
+        # DAO Treasury tools
+        "dao_treasury_list": TreasuryListTool(account_index),
         "dao_treasury_deploy": TreasuryDeployTool(account_index),
         "dao_treasury_deposit": TreasuryDepositTool(account_index),
         "dao_treasury_withdraw": TreasuryWithdrawTool(account_index),