feat: update agents functions (#20)

Author: fishonamos

.gitignore

@@ -19,6 +19,7 @@ npm-debug.log*
 yarn-debug.log*
 yarn-error.log*
 pnpm-debug.log*
+packages/sui-agent/src/examples/
 
 # local env files
 .env

apps/web/.env.example

@@ -3,6 +3,6 @@ PORT=2512
 
 # Atoma SDK Configuration
 # Get your bearer token from https://atoma.network
-ATOMASDK_BEARER_AUTH=your_bearer_token_here
+ATOMASDK_BEARER_AUTH=
 # Specify the Atoma chat model to use
 ATOMA_CHAT_COMPLETIONS_MODEL=meta-llama/Llama-3.3-70B-Instruct

packages/sui-agent/src/@types/interface.ts

@@ -119,6 +119,9 @@ export const COIN_SYNONYMS: Record<string, string> = {
   // AI and variants
   AI: 'AI',
   AICOIN: 'AI',
+  // WSB and variants
+  WSB: 'WSB',
+  WSBCOIN: 'WSB',
 } as const;
 
 // Mapping of coin symbols to their respective addresses on the Sui network
@@ -158,6 +161,7 @@ export const COIN_ADDRESSES = {
   SUIAI:
     '0x8f6808fcea9d5143e6d1577822fd49e783cb5ad2be042f4295e7fe2d5cb10b31::suiai::SUIAI',
   AI: '0x089de9a53ffd1f9252cf97e32e11b9a242f813e34227362b674b963468ec6620::ai::AI',
+  WSB: '0x4db126eac4fa99207e98db61d968477021fdeae153de3b244bcfbdc468ef0722::wsb::WSB',
 } as const;
 
 // Information about a liquidity pool

packages/sui-agent/src/agents/SuiAgent.ts

@@ -1,5 +1,5 @@
 import { IntentAgentResponse } from '../@types/interface';
-import Tools from '../tools/aftermath';
+import Tools from '../utils/tools';
 import { registerAllTools } from './ToolRegistry';
 import Utils from '../utils';
 import intent_agent_prompt from '../prompts/intent_agent_prompt';

packages/sui-agent/src/agents/ToolRegistry.ts

@@ -1,4 +1,4 @@
-import Tools from '../tools/aftermath';
+import Tools from '../utils/tools';
 import { getCoinPrice, coinsToPrice } from '../tools/aftermath/PriceTool';
 import { getTokenAPR } from '../tools/aftermath/apr';
 import {

packages/sui-agent/src/tools/aftermath/TradeTool.ts

@@ -22,21 +22,38 @@ async function findTokensInPool(
   console.log('Available tokens:', availableTokens);
   console.log('Looking for tokens:', { coinInType, coinOutType });
 
-  // Find matching tokens (case-insensitive and partial match)
+  // Find matching tokens (case-insensitive and handle both full addresses and short names)
   const findToken = (searchToken: string) => {
-    return availableTokens.find((token) =>
-      token.toLowerCase().includes(searchToken.toLowerCase()),
+    // If it's a full address, try exact match first
+    const exactMatch = availableTokens.find(
+      (token) => token.toLowerCase() === searchToken.toLowerCase(),
     );
+    if (exactMatch) return exactMatch;
+
+    // For short names like 'SUI' or 'WSB', look for them in the token addresses
+    return availableTokens.find((token) => {
+      // Normalize the token string to handle malformed addresses
+      const normalizedToken = token.replace(/:{2,}/g, '::').toLowerCase();
+      const parts = normalizedToken.split('::');
+
+      // Get the token name (last part) and clean it
+      const tokenName = parts[parts.length - 1];
+      const searchName = searchToken.toLowerCase();
+
+      return (
+        tokenName === searchName ||
+        tokenName.includes(searchName) ||
+        searchName.includes(tokenName)
+      );
+    });
   };
 
   const matchedCoinIn = findToken(coinInType);
   const matchedCoinOut = findToken(coinOutType);
 
   if (!matchedCoinIn || !matchedCoinOut) {
     throw new Error(
-      `Tokens not found in pool. Available tokens: ${availableTokens.join(
-        ', ',
-      )}`,
+      `Tokens not found in pool. Available tokens: ${availableTokens.join(', ')}`,
     );
   }
 

packages/sui-agent/src/transactions/TransactionTool.ts

@@ -61,6 +61,9 @@ export async function buildTransferTx(
 ): Promise<TransactionBlock> {
   const tx = new TransactionBlock();
 
+  // Set gas budget
+  tx.setGasBudget(2000000);
+
   // Get coins owned by sender
   const coins = await client.getCoins({
     owner: fromAddress,

packages/sui-agent/src/utils/index.ts

@@ -1,7 +1,7 @@
 import { randomUUID } from 'crypto';
 import { AtomaSDK } from 'atoma-sdk';
 import { atomaChat } from '../config/atoma';
-import Tools from '../tools/aftermath';
+import Tools from './tools';
 import { ToolArgument } from '../@types/interface';
 
 /**

…s/sui-agent/src/tools/aftermath/index.ts packages/sui-agent/src/utils/tools.tspackages/sui-agent/src/tools/aftermath/index.ts renamed to packages/sui-agent/src/utils/tools.ts packages/sui-agent/src/tools/aftermath/index.ts renamed to packages/sui-agent/src/utils/tools.ts

@@ -1,5 +1,5 @@
-import { Tool, toolResponse } from '../../@types/interface';
-import { atomaChat } from '../../config/atoma';
+import { Tool, toolResponse } from '../@types/interface';
+import { atomaChat } from '../config/atoma';
 import { AtomaSDK } from 'atoma-sdk';
 
 /**