Merge pull request #24 from FuelLabs/dp/get-balance

fix `agent.execute` and add balance read queries

src/FuelAgent.ts

@@ -11,6 +11,7 @@ import {
 } from './transfers/transfers.js';
 import { createAgent } from './agent.js';
 import { AgentExecutor } from 'langchain/agents';
+import { getOwnBalance, type GetOwnBalanceParams } from './read/balance.js';
 
 interface FuelAgentConfig {
   walletPrivateKey: string;
@@ -34,7 +35,7 @@ export class FuelAgent {
       throw new Error('Fuel wallet private key is required.');
     }
 
-    this.agentExecutor = createAgent(this.openAIKey);
+    this.agentExecutor = createAgent(this.openAIKey, this);
   }
 
   getCredentials() {
@@ -71,4 +72,8 @@ export class FuelAgent {
   async addLiquidity(params: AddLiquidityParams) {
     return await addLiquidity(params, this.walletPrivateKey);
   }
+
+  async getOwnBalance(params: GetOwnBalanceParams) {
+    return await getOwnBalance(params, this.walletPrivateKey);
+  }
 }

src/agent.ts

@@ -1,7 +1,7 @@
 import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { ChatOpenAI } from '@langchain/openai';
 import { createToolCallingAgent, AgentExecutor } from 'langchain/agents';
-import { tools } from './tools.js';
+import { createTools } from './tools.js';
 
 export const prompt = ChatPromptTemplate.fromMessages([
   [
@@ -19,12 +19,17 @@ export const prompt = ChatPromptTemplate.fromMessages([
   ['placeholder', '{agent_scratchpad}'],
 ]);
 
-export const createAgent = (openAIKey: string) => {
+export const createAgent = (
+  openAIKey: string, 
+  fuelAgent: { getCredentials: () => { walletPrivateKey: string } }
+) => {
   const model = new ChatOpenAI({
-    modelName: 'gpt-4o',
+    modelName: 'gpt-4',
     apiKey: openAIKey,
   });
 
+  const tools = createTools(fuelAgent);
+
   const agent = createToolCallingAgent({
     llm: model,
     tools,

src/mira/addLiquidity.ts

@@ -79,10 +79,7 @@ export const addLiquidity = async (
   if (!amount1InWei) {
     throw new Error('Failed to calculate amount1');
   }
-
-  console.log('Amount0 (Wei):', amount0InWei.toString());
-  console.log('Estimated Amount1 (Wei):', amount1InWei.toString());
-
+
   // Calculate minimum amounts with slippage
   const minAmount0 = amount0InWei
     .mul(bn(100 - Math.floor((params.slippage || DEFAULT_SLIPPAGE) * 100)))
@@ -91,9 +88,6 @@ export const addLiquidity = async (
     .mul(bn(100 - Math.floor((params.slippage || DEFAULT_SLIPPAGE) * 100)))
     .div(bn(100));
 
-  console.log('Min Amount0 (Wei):', minAmount0.toString());
-  console.log('Min Amount1 (Wei):', minAmount1.toString());
-
   const req = await miraAmm.addLiquidity(
     poolId,
     amount0InWei,

src/mira/swap.ts

@@ -90,10 +90,7 @@ export const swapExactInput = async (
   if (!amountOutWei) {
     throw new Error('Failed to calculate output amount');
   }
-
-  console.log('Amount In (Wei):', amountInWei.toString());
-  console.log('Estimated Amount Out (Wei):', amountOutWei.toString());
-
+
   const minAmountOut = amountOutWei
     .mul(bn(100 - Math.floor((params.slippage || DEFAULT_SLIPPAGE) * 100)))
     .div(bn(100));

src/read/balance.ts

@@ -0,0 +1,52 @@
+import { Provider } from 'fuels';
+import { getAllVerifiedFuelAssets } from '../utils/assets.js';
+import { setupWallet } from '../utils/setup.js';
+
+export type GetOwnBalanceParams = {
+  symbol: string;
+};
+
+export const getOwnBalance = async (
+  params: GetOwnBalanceParams,
+  privateKey: string,
+) => {
+  const { wallet } = await setupWallet(privateKey);
+
+  const allAssets = await getAllVerifiedFuelAssets();
+  const asset = allAssets.find((asset) => asset.symbol === params.symbol);
+
+  if (!asset) {
+    throw new Error(`Asset ${params.symbol} not found`);
+  }
+
+  const balance = await wallet.getBalance(asset.assetId);
+
+  return `Your ${params.symbol} balance is ${balance.formatUnits(asset.decimals)} ${params.symbol}`;
+};
+
+// ===============================
+
+export type GetBalanceParams = {
+  walletAddress: string;
+  assetSymbol: string;
+};
+
+export const getBalance = async (params: GetBalanceParams) => {
+  const provider = await Provider.create(
+    'https://mainnet.fuel.network/v1/graphql',
+  );
+
+  const allAssets = await getAllVerifiedFuelAssets();
+  const asset = allAssets.find((asset) => asset.symbol === params.assetSymbol);
+
+  if (!asset) {
+    throw new Error(`Asset ${params.assetSymbol} not found`);
+  }
+
+  const balance = await provider.getBalance(
+    params.walletAddress,
+    asset.assetId,
+  );
+
+  return `The ${params.assetSymbol} balance of ${params.walletAddress} is ${balance.formatUnits(asset.decimals)} ${params.assetSymbol}`;
+};

src/tools.ts

@@ -1,86 +1,127 @@
 import { tool } from '@langchain/core/tools';
 import { z } from 'zod';
+
+// Import functions
 import { swapExactInput } from './mira/swap.js';
 import { transfer } from './transfers/transfers.js';
 import { supplyCollateral } from './swaylend/supply.js';
 import { borrowAsset } from './swaylend/borrow.js';
 import { addLiquidity } from './mira/addLiquidity.js';
+import { getBalance, getOwnBalance } from './read/balance.js';
+
+// Types
+type FuelAgentInterface = {
+  getCredentials: () => { walletPrivateKey: string };
+};
 
 /**
- * Wrapper functions that only take the params argument
- * @param fn - The function to wrap.
- * @returns A new function that takes only the params argument.
+ * Wraps a function to inject the wallet private key from the agent
+ * @param fn - The function to wrap
+ * @param agent - The FuelAgent instance containing credentials
  */
-const wrapWithoutPrivateKey = <T>(
+const withWalletKey = <T>(
   fn: (params: T, privateKey: string) => Promise<any>,
+  agent: FuelAgentInterface,
 ) => {
-  return (params: T) => fn(params, '');
+  return (params: T) => fn(params, agent.getCredentials().walletPrivateKey);
 };
 
-export const transferTool = tool(wrapWithoutPrivateKey(transfer), {
-  name: 'fuel_transfer',
-  description: 'Transfer any verified Fuel asset to another wallet',
-  schema: z.object({
-    to: z.string().describe('The wallet address to transfer to'),
-    amount: z.string().describe('The amount to transfer'),
-    symbol: z.string().describe('The asset symbol to transfer. eg. USDC, ETH'),
-  }),
+// Schema definitions
+const transferSchema = z.object({
+  to: z.string().describe('The wallet address to transfer to'),
+  amount: z.string().describe('The amount to transfer'),
+  symbol: z.string().describe('The asset symbol to transfer. eg. USDC, ETH'),
 });
 
-export const swapExactInputTool = tool(wrapWithoutPrivateKey(swapExactInput), {
-  name: 'swap_exact_input',
-  description: 'Swap exact input on Mira',
-  schema: z.object({
-    amount: z.string().describe('The amount to swap'),
-    fromSymbol: z
-      .string()
-      .describe('The asset symbol to swap from. eg. USDC, ETH'),
-    toSymbol: z.string().describe('The asset symbol to swap to. eg. USDC, ETH'),
-    slippage: z
-      .number()
-      .optional()
-      .describe('Slippage tolerance (default: 0.01 for 1%)'),
-  }),
+const swapSchema = z.object({
+  amount: z.string().describe('The amount to swap'),
+  fromSymbol: z
+    .string()
+    .describe('The asset symbol to swap from. eg. USDC, ETH'),
+  toSymbol: z.string().describe('The asset symbol to swap to. eg. USDC, ETH'),
+  slippage: z
+    .number()
+    .optional()
+    .describe('Slippage tolerance (default: 0.01 for 1%)'),
+});
+
+const supplyCollateralSchema = z.object({
+  amount: z.string().describe('The amount to lend'),
+  symbol: z.string().describe('The asset symbol to lend. eg. USDC, ETH'),
+});
+
+const borrowAssetSchema = z.object({
+  amount: z.string().describe('The amount to borrow'),
+});
+
+const addLiquiditySchema = z.object({
+  amount0: z.string().describe('The amount of the first asset to add'),
+  asset0Symbol: z.string().describe('The symbol of the first asset'),
+  asset1Symbol: z.string().describe('The symbol of the second asset'),
+  slippage: z
+    .number()
+    .optional()
+    .describe('Slippage tolerance (default: 0.01 for 1%)'),
+});
+
+const getOwnBalanceSchema = z.object({
+  symbol: z
+    .string()
+    .describe('The asset symbol to get the balance of. eg. USDC, ETH'),
+});
+
+const getBalanceSchema = z.object({
+  walletAddress: z
+    .string()
+    .describe('The wallet address to get the balance of'),
+  assetSymbol: z
+    .string()
+    .describe('The asset symbol to get the balance of. eg. USDC, ETH'),
 });
 
-export const supplyCollateralTool = tool(
-  wrapWithoutPrivateKey(supplyCollateral),
-  {
+/**
+ * Creates and returns all tools with injected agent credentials
+ */
+export const createTools = (agent: FuelAgentInterface) => [
+  tool(withWalletKey(transfer, agent), {
+    name: 'fuel_transfer',
+    description: 'Transfer any verified Fuel asset to another wallet',
+    schema: transferSchema,
+  }),
+
+  tool(withWalletKey(swapExactInput, agent), {
+    name: 'swap_exact_input',
+    description: 'Swap exact input on Mira',
+    schema: swapSchema,
+  }),
+
+  tool(withWalletKey(supplyCollateral, agent), {
     name: 'supply_collateral',
     description: 'Supply collateral on swaylend',
-    schema: z.object({
-      amount: z.string().describe('The amount to lend'),
-      symbol: z.string().describe('The asset symbol to lend. eg. USDC, ETH'),
-    }),
-  },
-);
-
-export const borrowAssetTool = tool(wrapWithoutPrivateKey(borrowAsset), {
-  name: 'borrow_asset',
-  description: 'Borrow asset on swaylend',
-  schema: z.object({
-    amount: z.string().describe('The amount to borrow'),
+    schema: supplyCollateralSchema,
   }),
-});
 
-export const addLiquidityTool = tool(wrapWithoutPrivateKey(addLiquidity), {
-  name: 'add_liquidity',
-  description: 'Add liquidity to a Mira pool',
-  schema: z.object({
-    amount0: z.string().describe('The amount of the first asset to add'),
-    asset0Symbol: z.string().describe('The symbol of the first asset'),
-    asset1Symbol: z.string().describe('The symbol of the second asset'),
-    slippage: z
-      .number()
-      .optional()
-      .describe('Slippage tolerance (default: 0.01 for 1%)'),
+  tool(withWalletKey(borrowAsset, agent), {
+    name: 'borrow_asset',
+    description: 'Borrow asset on swaylend',
+    schema: borrowAssetSchema,
+  }),
+
+  tool(withWalletKey(addLiquidity, agent), {
+    name: 'add_liquidity',
+    description: 'Add liquidity to a Mira pool',
+    schema: addLiquiditySchema,
   }),
-});
 
-export const tools = [
-  transferTool,
-  swapExactInputTool,
-  supplyCollateralTool,
-  borrowAssetTool,
-  addLiquidityTool,
+  tool(withWalletKey(getOwnBalance, agent), {
+    name: 'get_own_balance',
+    description: 'Get the balance of an asset in your wallet',
+    schema: getOwnBalanceSchema,
+  }),
+
+  tool(getBalance, {
+    name: 'get_balance',
+    description: 'Get the balance of an asset for a given wallet address',
+    schema: getBalanceSchema,
+  }),
 ];

test/balance.test.ts

@@ -0,0 +1,33 @@
+import { beforeEach, test } from 'vitest';
+import { createTestAgent } from './setup.js';
+import type { FuelAgent } from '../src/FuelAgent.js';
+
+let agent: FuelAgent;
+
+beforeEach(() => {
+  agent = createTestAgent();
+});
+
+test(
+  'get own balance',
+  async () => {
+    const balance = await agent.execute('Get my USDC balance');
+    console.log(balance);
+  },
+  {
+    timeout: 500000,
+  },
+);
+
+test(
+  'get balance of a wallet',
+  async () => {
+    const balance = await agent.execute(
+      'Get the ETH balance of 0x8F8afB12402C9a4bD9678Bec363E51360142f8443FB171655eEd55dB298828D1',
+    );
+    console.log(balance);
+  },
+  {
+    timeout: 500000,
+  },
+);