Skip to content

Commit

Permalink
halfway finished testing e2e
Browse files Browse the repository at this point in the history
  • Loading branch information
danXyu committed Jan 9, 2025
1 parent fd74a80 commit ac9fe0e
Show file tree
Hide file tree
Showing 37 changed files with 1,621 additions and 157 deletions.
18 changes: 15 additions & 3 deletions submodules/moragents_dockers/agents/src/agents/token_swap/agent.py
Original file line number Diff line number Diff line change
@@ -1,22 +1,30 @@
import logging
import requests
from dataclasses import dataclass
from src.agents.token_swap import tools
from src.models.core import ChatRequest, AgentResponse
from src.agents.agent_core.agent import AgentCore
from langchain.schema import HumanMessage, SystemMessage
from src.stores.key_manager import key_manager_instance
from src.agents.token_swap.config import Config

logger = logging.getLogger(__name__)


@dataclass
class TokenSwapContext:
chain_id: str = None
wallet_address: str = None


class TokenSwapAgent(AgentCore):
"""Agent for handling token swap operations."""

def __init__(self, config, llm, embeddings):
super().__init__(config, llm, embeddings)
self.tools_provided = tools.get_tools()
self.tool_bound_llm = self.llm.bind_tools(self.tools_provided)
self.context = []
self.context = TokenSwapContext()

def _api_request_url(self, method_name, query_params, chain_id):
base_url = self.config.APIBASEURL + str(chain_id)
Expand Down Expand Up @@ -53,6 +61,10 @@ async def _process_request(self, request: ChatRequest) -> AgentResponse:
content="To help you with token swaps, I need your 1inch API key. Please set it up in Settings first."
)

# Store request context
self.context.chain_id = request.chain_id
self.context.wallet_address = request.wallet_address

messages = [
SystemMessage(
content=(
Expand Down Expand Up @@ -93,8 +105,8 @@ async def _execute_tool(self, func_name: str, args: dict) -> AgentResponse:
args["token1"],
args["token2"],
float(args["value"]),
self.config.chain_id,
self.config.wallet_address,
self.context.chain_id,
self.context.wallet_address,
)

# Return an action_required response with swap details
Expand Down
61 changes: 42 additions & 19 deletions submodules/moragents_dockers/agents/src/agents/token_swap/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,11 @@

import requests
from src.agents.token_swap.config import Config
from src.stores import key_manager_instance
from web3 import Web3

logger = logging.getLogger(__name__)


class InsufficientFundsError(Exception):
pass
Expand All @@ -18,22 +21,40 @@ class SwapNotPossibleError(Exception):
pass


def search_tokens(query, chain_id, limit=1, ignore_listed="false"):
def get_headers() -> dict[str, str]:
"""Get headers for 1inch API requests with optional API key override"""
oneinch_keys = key_manager_instance.get_oneinch_keys()
headers = {
"Authorization": f"Bearer {oneinch_keys.api_key}",
"accept": "application/json",
}
return headers


def search_tokens(
query: str,
chain_id: int,
limit: int = 1,
ignore_listed: str = "false",
) -> dict | None:
logger.info(f"Searching tokens - Query: {query}, Chain ID: {chain_id}")
endpoint = f"/v1.2/{chain_id}/search"
params = {"query": query, "limit": limit, "ignore_listed": ignore_listed}
response = requests.get(Config.INCH_URL + endpoint, params=params, headers=Config.HEADERS)
params = {"query": str(query), "limit": str(limit), "ignore_listed": str(ignore_listed)}

response = requests.get(Config.INCH_URL + endpoint, params=params, headers=get_headers())
logger.info(f"Search tokens response status: {response.status_code}")
if response.status_code == 200:
return response.json()
result = response.json()
logger.info(f"Found tokens: {result}")
return result
else:
logging.error(f"Failed to search tokens. Status code: {response.status_code}")
logger.error(f"Failed to search tokens. Status code: {response.status_code}, Response: {response.text}")
return None


def get_token_balance(web3: Web3, wallet_address: str, token_address: str, abi: list) -> int:
"""Get the balance of an ERC-20 token for a given wallet address."""
if (
not token_address
): # If no token address is provided, assume checking ETH or native token balance
if not token_address: # If no token address is provided, assume checking ETH or native token balance
return web3.eth.get_balance(web3.to_checksum_address(wallet_address))
else:
contract = web3.eth.contract(address=web3.to_checksum_address(token_address), abi=abi)
Expand Down Expand Up @@ -83,30 +104,34 @@ def validate_swap(web3: Web3, token1, token2, chain_id, amount, wallet_address):
raise TokenNotFoundError(f"Token {token2} not found.")

# Check if the user has sufficient balance for the swap
if t1_bal < smallest_amount:
raise InsufficientFundsError(f"Insufficient funds to perform the swap.")
# if t1_bal < smallest_amount:
# raise InsufficientFundsError(f"Insufficient funds to perform the swap.")

return t1[0]["address"], t1[0]["symbol"], t2[0]["address"], t2[0]["symbol"]


def get_quote(token1, token2, amount_in_wei, chain_id):
logger.info(f"Getting quote - Token1: {token1}, Token2: {token2}, Amount: {amount_in_wei}, Chain ID: {chain_id}")
endpoint = f"/v6.0/{chain_id}/quote"
params = {"src": token1, "dst": token2, "amount": int(amount_in_wei)}
response = requests.get(Config.QUOTE_URL + endpoint, params=params, headers=Config.HEADERS)
logger.debug(f"Quote request - URL: {Config.QUOTE_URL + endpoint}, Params: {params}")

response = requests.get(Config.QUOTE_URL + endpoint, params=params, headers=get_headers())
logger.info(f"Quote response status: {response.status_code}")
if response.status_code == 200:
return response.json()
result = response.json()
logger.info(f"Quote received: {result}")
return result
else:
logging.error(f"Failed to get quote. Status code: {response.status_code}")
logger.error(f"Failed to get quote. Status code: {response.status_code}, Response: {response.text}")
return None


def get_token_decimals(web3: Web3, token_address: str) -> int:
if not token_address:
return 18 # Assuming 18 decimals for the native gas token
else:
contract = web3.eth.contract(
address=Web3.to_checksum_address(token_address), abi=Config.ERC20_ABI
)
contract = web3.eth.contract(address=Web3.to_checksum_address(token_address), abi=Config.ERC20_ABI)
return contract.functions.decimals().call()


Expand Down Expand Up @@ -135,9 +160,7 @@ def swap_coins(token1, token2, amount, chain_id, wallet_address):
t2_address = "" if t2_a == Config.INCH_NATIVE_TOKEN_ADDRESS else t2_a
t2_quote = convert_to_readable_unit(web3, int(price), t2_address)
else:
raise SwapNotPossibleError(
"Failed to generate a quote. Please ensure you're on the correct network."
)
raise SwapNotPossibleError("Failed to generate a quote. Please ensure you're on the correct network.")

return {
"dst": t2_id,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import React, { FC } from "react";
import { Box, Flex, Image, Text } from "@chakra-ui/react";
import styles from "./index.module.css";
import styles from "./ImageDisplayMessage.module.css";

type ImageDisplayProps = {
imageMetadata: {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import React, { FC } from "react";
import { ClaimForm } from "../ClaimForm/ClaimForm";
import { ClaimForm } from "./ClaimForm";
import { ClaimMessagePayload } from "@/services/types";

type ClaimMessageProps = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ import {
} from "@chakra-ui/react";
import { useAccount, useChainId, useReadContract } from "wagmi";
import { erc20Abi, zeroAddress } from "viem";
import { oneInchNativeToken, routerAddress } from "../../config";
import { oneInchNativeToken, routerAddress } from "../../../config";
import { InfoIcon } from "@chakra-ui/icons";

export type SwapMessageLike =
Expand Down
Original file line number Diff line number Diff line change
@@ -1,21 +1,20 @@
import React, { FC } from "react";
import { SwapForm } from "../SwapForm";
import { SwapMessagePayload } from "@/services/types";
import { Box, Text } from "@chakra-ui/react";
import { SwapMessagePayload } from "@/services/types";
import { SwapForm } from "@/components/Agents/Swaps/SwapForm";
import { useSwapTransaction } from "@/components/Agents/Swaps/useSwapTransaction";

type SwapMessageProps = {
isActive: boolean;
onCancelSwap: (fromAction: number) => void;
fromMessage: SwapMessagePayload | null;
onSubmitSwap: (swapTx: any) => void;
};

export const SwapMessage: FC<SwapMessageProps> = ({
isActive,
onCancelSwap,
fromMessage,
onSubmitSwap,
}) => {
const { handleSwap, handleCancel, isLoading } = useSwapTransaction();

if (!fromMessage) {
return (
<Box p={4} bg="red.100" color="red.800" borderRadius="md">
Expand All @@ -27,10 +26,10 @@ export const SwapMessage: FC<SwapMessageProps> = ({
return (
<SwapForm
isActive={isActive}
onCancelSwap={onCancelSwap}
fromMessage={fromMessage}
onSubmitApprove={() => {}} // Implement approve logic if needed
onSubmitSwap={onSubmitSwap}
onCancelSwap={handleCancel}
onSubmitSwap={handleSwap}
isLoading={isLoading}
/>
);
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
import { useState, useCallback } from "react";
import { useAccount, useChainId, useSendTransaction } from "wagmi";
import { sendSwapStatus } from "@/services/apiHooks";
import { getHttpClient, SWAP_STATUS } from "@/services/constants";
import { SwapTxPayloadType } from "@/services/types";

export const useSwapTransaction = () => {
const [isLoading, setIsLoading] = useState(false);
const [txHash, setTxHash] = useState("");
const { address } = useAccount();
const chainId = useChainId();
const { sendTransaction } = useSendTransaction();

const handleSwap = useCallback(
async (swapTx: SwapTxPayloadType) => {
if (!address) return;

setIsLoading(true);
setTxHash("");

try {
sendTransaction(
{
account: address,
data: (swapTx?.tx.data || "0x") as `0x${string}`,
to: (swapTx?.tx.to || "0x") as `0x${string}`,
value: BigInt(swapTx?.tx.value || "0"),
},
{
onSuccess: async (hash) => {
setTxHash(hash);
await sendSwapStatus(
getHttpClient(),
chainId,
address.toLowerCase(),
SWAP_STATUS.INIT,
hash,
0
);
},
onError: async (error) => {
console.error(`Error sending transaction: ${error}`);
await sendSwapStatus(
getHttpClient(),
chainId,
address.toLowerCase(),
SWAP_STATUS.FAIL,
"",
0
);
},
onSettled: () => setIsLoading(false),
}
);
} catch (error) {
setIsLoading(false);
console.error("Swap failed:", error);
}
},
[address, chainId, sendTransaction]
);

const handleCancel = useCallback(
async (fromAction: number) => {
if (!address) return;

setIsLoading(true);
try {
await sendSwapStatus(
getHttpClient(),
chainId,
address,
SWAP_STATUS.CANCELLED,
"",
fromAction
);
} catch (error) {
console.error(`Failed to cancel swap: ${error}`);
} finally {
setIsLoading(false);
}
},
[address, chainId]
);

return {
handleSwap,
handleCancel,
isLoading,
txHash,
};
};
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import {
} from "react-icons/fa";
import { postTweet, regenerateTweet } from "@/services/apiHooks";
import { getHttpClient } from "@/services/constants";
import styles from "./index.module.css";
import styles from "./TweetMessage.module.css";

type TweetProps = {
initialContent: string;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
import React from "react";
import {
VStack,
FormControl,
FormLabel,
Input,
Select,
Button,
} from "@chakra-ui/react";
import styles from "./index.module.css";
interface CreateWalletFormProps {
walletName: string;
selectedNetwork: string;
networks: string[];
onWalletNameChange: (value: string) => void;
onNetworkChange: (value: string) => void;
onSubmit: () => void;
}

export const CreateWalletForm: React.FC<CreateWalletFormProps> = ({
walletName,
selectedNetwork,
networks,
onWalletNameChange,
onNetworkChange,
onSubmit,
}) => (
<VStack spacing={4} className={styles.createWalletForm}>
<FormControl>
<FormLabel className={styles.formLabel}>Wallet Name</FormLabel>
<Input
className={styles.input}
placeholder="Enter wallet name"
value={walletName}
onChange={(e) => onWalletNameChange(e.target.value)}
/>
</FormControl>
<FormControl>
<FormLabel className={styles.formLabel}>Network</FormLabel>
<Select
className={styles.select}
value={selectedNetwork}
onChange={(e) => onNetworkChange(e.target.value)}
>
{networks.map((network) => (
<option key={network} value={network}>
{network}
</option>
))}
</Select>
</FormControl>
<Button className={styles.submitButton} onClick={onSubmit}>
Create Wallet
</Button>
</VStack>
);
Loading

0 comments on commit ac9fe0e

Please sign in to comment.