From 2e1c533b14fdca556bf940b43b16465680bd7286 Mon Sep 17 00:00:00 2001 From: TxCorpi0x <6095314+TxCorpi0x@users.noreply.github.com> Date: Mon, 3 Feb 2025 15:16:16 +0300 Subject: [PATCH 1/2] feat: enso skills async --- app/core/engine.py | 1 + skills/enso/__init__.py | 36 +++++++++------ skills/enso/base.py | 12 +++-- skills/enso/networks.py | 17 +++++-- skills/enso/prices.py | 19 +++++++- skills/enso/route.py | 58 ++++++++++++++--------- skills/enso/tokens.py | 23 +++++++-- skills/enso/wallet.py | 100 ++++++++++++++++++++++++++++++---------- 8 files changed, 192 insertions(+), 74 deletions(-) diff --git a/app/core/engine.py b/app/core/engine.py index 7b304932..20c3b9f1 100644 --- a/app/core/engine.py +++ b/app/core/engine.py @@ -192,6 +192,7 @@ async def initialize_agent(aid): agentkit.wallet if agentkit else None, config.rpc_base_mainnet, skill_store, + agent_store, aid, ) tools.append(s) diff --git a/skills/enso/__init__.py b/skills/enso/__init__.py index 8345f270..2c98c460 100644 --- a/skills/enso/__init__.py +++ b/skills/enso/__init__.py @@ -9,7 +9,7 @@ from skills.enso.route import EnsoRouteShortcut from skills.enso.tokens import EnsoGetTokens from skills.enso.wallet import ( - EnsoBroadcastWalletApprove, + EnsoWalletApprove, EnsoGetWalletApprovals, EnsoGetWalletBalances, ) @@ -20,8 +20,9 @@ def get_enso_skill( api_token: str, main_tokens: list[str], wallet: Wallet, - rpc_nodes: dict[str, str], - store: SkillStoreABC, + rpc_node: str, + skill_store: SkillStoreABC, + agent_store: SkillStoreABC, agent_id: str, ) -> EnsoBaseTool: if not api_token: @@ -31,7 +32,8 @@ def get_enso_skill( return EnsoGetNetworks( api_token=api_token, main_tokens=main_tokens, - store=store, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) @@ -39,7 +41,8 @@ def get_enso_skill( return EnsoGetTokens( api_token=api_token, main_tokens=main_tokens, - store=store, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) @@ -47,7 +50,8 @@ def get_enso_skill( return EnsoGetPrices( api_token=api_token, main_tokens=main_tokens, - store=store, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) @@ -58,7 +62,8 @@ def get_enso_skill( api_token=api_token, main_tokens=main_tokens, wallet=wallet, - store=store, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) @@ -69,31 +74,34 @@ def get_enso_skill( api_token=api_token, main_tokens=main_tokens, wallet=wallet, - store=store, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) if name == "wallet_approve": if not wallet: raise ValueError("Wallet is empty") - return EnsoBroadcastWalletApprove( + return EnsoWalletApprove( api_token=api_token, main_tokens=main_tokens, wallet=wallet, - rpc_nodes=rpc_nodes, - store=store, + rpc_node=rpc_node, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) - if name == "broadcast_route_shortcut": + if name == "route_shortcut": if not wallet: raise ValueError("Wallet is empty") return EnsoRouteShortcut( api_token=api_token, main_tokens=main_tokens, wallet=wallet, - rpc_nodes=rpc_nodes, - store=store, + rpc_node=rpc_node, + skill_store=skill_store, + agent_store=agent_store, agent_id=agent_id, ) diff --git a/skills/enso/base.py b/skills/enso/base.py index da05b52c..93ba2b68 100644 --- a/skills/enso/base.py +++ b/skills/enso/base.py @@ -3,6 +3,7 @@ from cdp import Wallet from pydantic import BaseModel, Field +from abstracts.agent import AgentStoreABC from abstracts.skill import IntentKitSkill, SkillStoreABC base_url = "https://api.enso.finance" @@ -15,11 +16,14 @@ class EnsoBaseTool(IntentKitSkill): api_token: str = Field(description="API token") main_tokens: list[str] = Field(description="Main supported tokens") wallet: Wallet | None = Field(None, description="The wallet of the agent") - rpc_nodes: dict[str, str] | None = Field( - None, description="RPC nodes for different networks" - ) + rpc_node: str | None = Field(None, description="RPC nodes for different networks") name: str = Field(description="The name of the tool") description: str = Field(description="A description of what the tool does") args_schema: Type[BaseModel] agent_id: str = Field(description="The ID of the agent") - store: SkillStoreABC = Field(description="The skill store for persisting data") + agent_store: AgentStoreABC = Field( + description="The agent store for persisting data" + ) + skill_store: SkillStoreABC = Field( + description="The skill store for persisting data" + ) diff --git a/skills/enso/networks.py b/skills/enso/networks.py index 92891c0f..e93356b4 100644 --- a/skills/enso/networks.py +++ b/skills/enso/networks.py @@ -45,6 +45,17 @@ class EnsoGetNetworks(EnsoBaseTool): args_schema: Type[BaseModel] = EnsoGetNetworksInput def _run(self) -> EnsoGetNetworksOutput: + """Run the tool to get all supported networks. + + Returns: + EnsoGetNetworksOutput: A structured output containing the result of the networks. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") + + async def _arun(self) -> EnsoGetNetworksOutput: """ Function to request the list of supported networks and their chain id and name. @@ -58,10 +69,10 @@ def _run(self) -> EnsoGetNetworksOutput: "Authorization": f"Bearer {self.api_token}", } - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: # Send the GET request - response = client.get(url, headers=headers) + response = await client.get(url, headers=headers) response.raise_for_status() # Parse the response JSON into the NetworkResponse model @@ -74,7 +85,7 @@ def _run(self) -> EnsoGetNetworksOutput: networks.append(network) networks_memory[network.id] = network.model_dump(exclude_none=True) - self.store.save_agent_skill_data( + await self.skill_store.save_agent_skill_data( self.agent_id, "enso_get_networks", "networks", diff --git a/skills/enso/prices.py b/skills/enso/prices.py index 1f434416..1fbaa34a 100644 --- a/skills/enso/prices.py +++ b/skills/enso/prices.py @@ -44,6 +44,21 @@ def _run( self, chainId: int = default_chain_id, address: str = "0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + ) -> EnsoGetPricesOutput: + """Run the tool to get the token price from the API. + + Returns: + EnsoGetPricesOutput: A structured output containing the result of token prices. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") + + async def _arun( + self, + chainId: int = default_chain_id, + address: str = "0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", ) -> EnsoGetPricesOutput: """ Asynchronous function to request the token price from the API. @@ -62,9 +77,9 @@ def _run( "Authorization": f"Bearer {self.api_token}", } - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: - response = client.get(url, headers=headers) + response = await client.get(url, headers=headers) response.raise_for_status() json_dict = response.json() diff --git a/skills/enso/route.py b/skills/enso/route.py index 3e791788..81845f24 100644 --- a/skills/enso/route.py +++ b/skills/enso/route.py @@ -156,7 +156,9 @@ class EnsoRouteShortcut(EnsoBaseTool): """ name: str = "enso_route_shortcut" - description: str = "This tool is used specifically for broadcasting a route transaction calldata to the network. It should only be used when the user explicitly requests to broadcast a route transaction with routeId." + description: str = ( + "This tool is used specifically for broadcasting a route transaction calldata to the network. It should only be used when the user explicitly requests to broadcast a route transaction with routeId." + ) args_schema: Type[BaseModel] = EnsoRouteShortcutInput def _run( @@ -170,6 +172,25 @@ def _run( """ Run the tool to get swap route information. + Returns: + EnsoRouteShortcutOutput: The response containing route shortcut information. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") + + async def _arun( + self, + amountIn: list[int], + tokenIn: list[str], + tokenOut: list[str], + chainId: int = default_chain_id, + broadcast_requested: bool | None = False, + ) -> EnsoRouteShortcutOutput: + """ + Run the tool to get swap route information. + Args: amountIn (list[int]): Amount of tokenIn to swap in wei, you should multiply user's requested value by token decimals. tokenIn (list[str]): Ethereum address of the token to swap or enter into a position from (For ETH, use 0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee). @@ -181,10 +202,10 @@ def _run( EnsoRouteShortcutOutput: The response containing route shortcut information. """ - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: network_name = None - networks = self.store.get_agent_skill_data( + networks = await self.skill_store.get_agent_skill_data( self.agent_id, "enso_get_networks", "networks" ) @@ -195,17 +216,15 @@ def _run( else None ) if network_name is None: - networks_list = ( - EnsoGetNetworks( - api_token=self.api_token, - main_tokens=self.main_tokens, - store=self.store, - agent_id=self.agent_id, - ) - .run(EnsoGetNetworksInput()) - .res - ) - for network in networks_list: + networks = await EnsoGetNetworks( + api_token=self.api_token, + main_tokens=self.main_tokens, + skill_store=self.skill_store, + agent_store=self.agent_store, + agent_id=self.agent_id, + ).arun(EnsoGetNetworksInput()) + + for network in networks.res: if network.id == chainId: network_name = network.name @@ -219,7 +238,7 @@ def _run( "Authorization": f"Bearer {self.api_token}", } - token_decimals = self.store.get_agent_skill_data( + token_decimals = await self.skill_store.get_agent_skill_data( self.agent_id, "enso_get_tokens", "decimals", @@ -252,7 +271,7 @@ def _run( params["fromAddress"] = self.wallet.addresses[0].address_id - response = client.get(url, headers=headers, params=params) + response = await client.get(url, headers=headers, params=params) response.raise_for_status() # Raise HTTPError for non-2xx responses json_dict = response.json() @@ -264,13 +283,8 @@ def _run( ) if broadcast_requested: - if not self.rpc_nodes.get(str(chainId)): - raise ToolException( - f"rpc node not found for chainId: {chainId}" - ) - contract = EvmContractWrapper( - self.rpc_nodes[str(chainId)], ABI_ROUTE, json_dict.get("tx") + self.rpc_node, ABI_ROUTE, json_dict.get("tx") ) fn, fn_args = contract.fn_and_args diff --git a/skills/enso/tokens.py b/skills/enso/tokens.py index f6d7ceaa..9016fd88 100644 --- a/skills/enso/tokens.py +++ b/skills/enso/tokens.py @@ -140,6 +140,21 @@ def _run( self, chainId: int = default_chain_id, protocolSlug: str | None = None, + ) -> EnsoGetTokensOutput: + """Run the tool to get the tokens and APYs from the API. + + Returns: + EnsoGetPricesOutput: A structured output containing the result of tokens and APYs. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") + + async def _arun( + self, + chainId: int = default_chain_id, + protocolSlug: str | None = None, ) -> EnsoGetTokensOutput: """Run the tool to get Tokens and APY. Args: @@ -165,13 +180,13 @@ def _run( params["page"] = 1 params["includeMetadata"] = "true" - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: - response = client.get(url, headers=headers, params=params) + response = await client.get(url, headers=headers, params=params) response.raise_for_status() json_dict = response.json() - token_decimals = self.store.get_agent_skill_data( + token_decimals = await self.skill_store.get_agent_skill_data( self.agent_id, "enso_get_tokens", "decimals", @@ -196,7 +211,7 @@ def _run( for u_token in token_response.underlyingTokens: token_decimals[u_token.address] = u_token.decimals - self.store.save_agent_skill_data( + await self.skill_store.save_agent_skill_data( self.agent_id, "enso_get_tokens", "decimals", diff --git a/skills/enso/wallet.py b/skills/enso/wallet.py index 9f030d62..1c434c71 100644 --- a/skills/enso/wallet.py +++ b/skills/enso/wallet.py @@ -60,7 +60,24 @@ class EnsoGetWalletBalances(EnsoBaseTool): ) args_schema: Type[BaseModel] = EnsoGetBalancesInput - def _run(self, chainId: int = default_chain_id) -> EnsoGetBalancesOutput: + def _run( + self, + chainId: int = default_chain_id, + ) -> EnsoGetBalancesOutput: + """Run the tool to get the token balances of a wallet. + + Returns: + EnsoGetPricesOutput: A structured output containing the result of token balances of a wallet. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") + + async def _arun( + self, + chainId: int = default_chain_id, + ) -> EnsoGetBalancesOutput: """ Run the tool to get token balances of a wallet. @@ -81,10 +98,10 @@ def _run(self, chainId: int = default_chain_id) -> EnsoGetBalancesOutput: params["eoaAddress"] = self.wallet.addresses[0].address_id params["useEoa"] = True - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: # Send the GET request - response = client.get(url, headers=headers, params=params) + response = await client.get(url, headers=headers, params=params) response.raise_for_status() # Map the response JSON into the WalletBalance model @@ -146,9 +163,28 @@ class EnsoGetWalletApprovals(EnsoBaseTool): description: str = ( "Retrieve token spend approvals for a wallet on a specified blockchain network." ) - args_schema: Type[BaseModel] = EnsoGetApprovalsInput + args_schema: Type[BaseModel] = EnsoGetApprovalsOutput + + def _run( + self, + chainId: int = default_chain_id, + **kwargs, + ) -> EnsoGetApprovalsOutput: + """Run the tool to get the token approvals for a wallet. + + Returns: + EnsoGetPricesOutput: A structured output containing the result of token approvals for a wallet. + + Raises: + Exception: If there's an error accessing the Enso API. + """ + raise NotImplementedError("Use _arun instead") - def _run(self, chainId: int = default_chain_id, **kwargs) -> EnsoGetApprovalsOutput: + async def _arun( + self, + chainId: int = default_chain_id, + **kwargs, + ) -> EnsoGetApprovalsOutput: """ Run the tool to get token approvals for a wallet. @@ -174,10 +210,10 @@ def _run(self, chainId: int = default_chain_id, **kwargs) -> EnsoGetApprovalsOut if kwargs.get("routingStrategy"): params.routingStrategy = kwargs["routingStrategy"] - with httpx.Client() as client: + async with httpx.AsyncClient() as client: try: # Send the GET request - response = client.get( + response = await client.get( url, headers=headers, params=params.model_dump(exclude_none=True) ) response.raise_for_status() @@ -200,7 +236,7 @@ def _run(self, chainId: int = default_chain_id, **kwargs) -> EnsoGetApprovalsOut raise ToolException(f"error from Enso API: {e}") from e -class EnsoBroadcastWalletApproveInput(BaseModel): +class EnsoWalletApproveInput(BaseModel): """ Input model for approve the wallet. """ @@ -215,7 +251,7 @@ class EnsoBroadcastWalletApproveInput(BaseModel): ) -class EnsoBroadcastWalletApproveOutput(BaseModel): +class EnsoWalletApproveOutput(BaseModel): """ Output model for approve token for the wallet. """ @@ -226,7 +262,7 @@ class EnsoBroadcastWalletApproveOutput(BaseModel): spender: str | None = Field(None, description="The spender address to approve") -class EnsoBroadcastWalletApproveArtifact(BaseModel): +class EnsoWalletApproveArtifact(BaseModel): """ Output model for approve token for the wallet. """ @@ -235,7 +271,7 @@ class EnsoBroadcastWalletApproveArtifact(BaseModel): txHash: str | None = Field(None, description="The transaction hash") -class EnsoBroadcastWalletApprove(EnsoBaseTool): +class EnsoWalletApprove(EnsoBaseTool): """ This tool is used specifically for broadcasting a ERC20 token spending approval transaction to the network. It should only be used when the user explicitly requests to broadcast an approval transaction with a specific amount for a certain token. @@ -249,23 +285,42 @@ class EnsoBroadcastWalletApprove(EnsoBaseTool): - Approving token spending grants another account permission to spend your tokens. Attributes: - name (str): Name of the tool, specifically "enso_broadcast_wallet_approve". + name (str): Name of the tool, specifically "enso_wallet_approve". description (str): Comprehensive description of the tool's purpose and functionality. args_schema (Type[BaseModel]): Schema for input arguments, specifying expected parameters. """ - name: str = "enso_broadcast_wallet_approve" - description: str = "This tool is used specifically for broadcasting a ERC20 token spending approval transaction to the network. It should only be used when the user explicitly requests to broadcast an approval transaction with a specific amount for a certain token." - args_schema: Type[BaseModel] = EnsoBroadcastWalletApproveInput + name: str = "enso_wallet_approve" + description: str = ( + "This tool is used specifically for broadcasting a ERC20 token spending approval transaction to the network. It should only be used when the user explicitly requests to broadcast an approval transaction with a specific amount for a certain token." + ) + args_schema: Type[BaseModel] = EnsoWalletApproveInput response_format: str = "content_and_artifact" + # def _run( + # self, + # tokenAddress: str, + # amount: int, + # chainId: int = default_chain_id, + # **kwargs, + # ) -> Tuple[EnsoBroadcastWalletApproveOutput, EnsoBroadcastWalletApproveArtifact]: + # """Run the tool to approve enso router for a wallet. + + # Returns: + # Tuple[EnsoBroadcastWalletApproveOutput, EnsoBroadcastWalletApproveArtifact]: A structured output containing the result of token approval. + + # Raises: + # Exception: If there's an error accessing the Enso API. + # """ + # raise NotImplementedError("Use _arun instead") + def _run( self, tokenAddress: str, amount: int, chainId: int = default_chain_id, **kwargs, - ) -> Tuple[EnsoBroadcastWalletApproveOutput, EnsoBroadcastWalletApproveArtifact]: + ) -> Tuple[EnsoWalletApproveOutput, EnsoWalletApproveArtifact]: """ Run the tool to approve enso router for a wallet. @@ -287,7 +342,7 @@ def _run( from_address = self.wallet.addresses[0].address_id - params = EnsoBroadcastWalletApproveInput( + params = EnsoWalletApproveInput( tokenAddress=tokenAddress, amount=amount, chainId=chainId, @@ -308,15 +363,10 @@ def _run( # Map the response JSON into the WalletApproveTransaction model json_dict = response.json() - content = EnsoBroadcastWalletApproveOutput(**json_dict) - artifact = EnsoBroadcastWalletApproveArtifact(**json_dict) + content = EnsoWalletApproveOutput(**json_dict) + artifact = EnsoWalletApproveArtifact(**json_dict) - if not self.rpc_nodes.get(str(chainId)): - raise ToolException(f"rpc node not found for chainId: {chainId}") - - contract = EvmContractWrapper( - self.rpc_nodes[str(chainId)], ABI_ERC20, artifact.tx - ) + contract = EvmContractWrapper(self.rpc_node, ABI_ERC20, artifact.tx) fn, fn_args = contract.fn_and_args fn_args["value"] = str(fn_args["value"]) From c9d3d9dcac702f4415610380e543e3143b982d8f Mon Sep 17 00:00:00 2001 From: TxCorpi0x <6095314+TxCorpi0x@users.noreply.github.com> Date: Mon, 3 Feb 2025 15:52:10 +0300 Subject: [PATCH 2/2] lint: ruff --- skills/enso/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skills/enso/__init__.py b/skills/enso/__init__.py index 2c98c460..dbd012e8 100644 --- a/skills/enso/__init__.py +++ b/skills/enso/__init__.py @@ -9,9 +9,9 @@ from skills.enso.route import EnsoRouteShortcut from skills.enso.tokens import EnsoGetTokens from skills.enso.wallet import ( - EnsoWalletApprove, EnsoGetWalletApprovals, EnsoGetWalletBalances, + EnsoWalletApprove, )