Merge pull request #26 from lakshyaag/lakshya/model-selection

feat: add environment configuration and support for multiple AI models

.env.example

@@ -0,0 +1,5 @@
+FUEL_WALLET_PRIVATE_KEY=
+
+OPENAI_API_KEY=
+ANTHROPIC_API_KEY=
+GOOGLE_GEMINI_API_KEY=

package.json

@@ -36,7 +36,9 @@
     "vitest": "^2.0.5"
   },
   "dependencies": {
+    "@langchain/anthropic": "^0.3.11",
     "@langchain/core": "^0.3.23",
+    "@langchain/google-genai": "^0.1.6",
     "@langchain/openai": "^0.3.14",
     "@pythnetwork/hermes-client": "^1.2.0",
     "@pythnetwork/pyth-fuel-js": "^1.0.7",

src/FuelAgent.ts

@@ -10,38 +10,52 @@ import {
   type TransferParams,
 } from './transfers/transfers.js';
 import { createAgent } from './agent.js';
-import { AgentExecutor } from 'langchain/agents';
+import type { AgentExecutor } from 'langchain/agents';
 import { getOwnBalance, type GetOwnBalanceParams } from './read/balance.js';
+import type { modelMapping } from './utils/models.js';
 
-interface FuelAgentConfig {
+export interface FuelAgentConfig {
   walletPrivateKey: string;
-  openAIKey: string;
+  model: keyof typeof modelMapping;
+  openAiApiKey?: string;
+  anthropicApiKey?: string;
+  googleGeminiApiKey?: string;
 }
 
-export class FuelAgent {
+export class FuelAgent {  
   private walletPrivateKey: string;
-  private openAIKey: string;
   private agentExecutor: AgentExecutor;
+  private model: keyof typeof modelMapping;
+  private openAiApiKey?: string;
+  private anthropicApiKey?: string;
+  private googleGeminiApiKey?: string;
 
   constructor(config: FuelAgentConfig) {
     this.walletPrivateKey = config.walletPrivateKey;
-    this.openAIKey = config.openAIKey;
-
-    if (!this.openAIKey) {
-      throw new Error('OpenAI API key is required.');
-    }
+    this.model = config.model;
+    this.openAiApiKey = config.openAiApiKey;
+    this.anthropicApiKey = config.anthropicApiKey;
+    this.googleGeminiApiKey = config.googleGeminiApiKey;
 
     if (!this.walletPrivateKey) {
       throw new Error('Fuel wallet private key is required.');
     }
 
-    this.agentExecutor = createAgent(this.openAIKey, this);
+    this.agentExecutor = createAgent(
+      this,
+      this.model,
+      this.openAiApiKey,
+      this.anthropicApiKey,
+      this.googleGeminiApiKey,
+    );
   }
 
   getCredentials() {
     return {
       walletPrivateKey: this.walletPrivateKey,
-      openAIKey: this.openAIKey,
+      openAiApiKey: this.openAiApiKey || '',
+      anthropicApiKey: this.anthropicApiKey || '',
+      googleGeminiApiKey: this.googleGeminiApiKey || '',
     };
   }
 

src/agent.ts

@@ -2,36 +2,75 @@ import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { ChatOpenAI } from '@langchain/openai';
 import { createToolCallingAgent, AgentExecutor } from 'langchain/agents';
 import { createTools } from './tools.js';
+import { modelMapping } from './utils/models.js';
+import { ChatAnthropic } from '@langchain/anthropic';
+import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
+import { SystemMessage } from '@langchain/core/messages';
+
+const systemMessage = new SystemMessage(
+  ` You are an AI agent on Fuel network capable of executing all kinds of transactions and interacting with the Fuel blockchain.
+    You are able to execute transactions on behalf of the user.
+
+    Always return the response in the following format:
+    The transaction was successful/failed. The explorer link is: https://app.fuel.network/tx/0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef/simple
+  `,
+);
 
 export const prompt = ChatPromptTemplate.fromMessages([
-  [
-    'system',
-    'You are an AI agent on Fuel network capable of executing all kinds of transactions and interacting with the Fuel blockchain.',
-  ],
-  [
-    'system',
-    `Always return the response in the following format:
-      The transaction was successful/failed. The explorer link is: https://app.fuel.network/tx/0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef/simple
-    `,
-  ],
+  systemMessage,
   ['placeholder', '{chat_history}'],
   ['human', '{input}'],
   ['placeholder', '{agent_scratchpad}'],
 ]);
 
 export const createAgent = (
-  openAIKey: string, 
-  fuelAgent: { getCredentials: () => { walletPrivateKey: string } }
+  fuelAgent: { getCredentials: () => { walletPrivateKey: string } },
+  modelName: keyof typeof modelMapping,
+  openAiApiKey?: string,
+  anthropicApiKey?: string,
+  googleGeminiApiKey?: string,
 ) => {
-  const model = new ChatOpenAI({
-    modelName: 'gpt-4',
-    apiKey: openAIKey,
-  });
+  const model = () => {
+    if (modelMapping[modelName] === 'openai') {
+      if (!openAiApiKey) {
+        throw new Error('OpenAI API key is required');
+      }
+      return new ChatOpenAI({
+        modelName: modelName,
+        apiKey: openAiApiKey,
+      });
+    }
+    if (modelMapping[modelName] === 'anthropic') {
+      if (!anthropicApiKey) {
+        throw new Error('Anthropic API key is required');
+      }
+      return new ChatAnthropic({
+        modelName: modelName,
+        anthropicApiKey: anthropicApiKey,
+      });
+    }
+    if (modelMapping[modelName] === 'gemini') {
+      if (!googleGeminiApiKey) {
+        throw new Error('Google Gemini API key is required');
+      }
+      return new ChatGoogleGenerativeAI({
+        modelName: modelName,
+        apiKey: googleGeminiApiKey,
+        convertSystemMessageToHumanContent: true,
+      });
+    }
+  };
+
+  const selectedModel = model();
+
+  if (!selectedModel) {
+    throw new Error('Error initializing model');
+  }
 
   const tools = createTools(fuelAgent);
 
   const agent = createToolCallingAgent({
-    llm: model,
+    llm: selectedModel,
     tools,
     prompt,
   });

src/mira/addLiquidity.ts

@@ -79,7 +79,7 @@ export const addLiquidity = async (
   if (!amount1InWei) {
     throw new Error('Failed to calculate amount1');
   }
-  
+
   // Calculate minimum amounts with slippage
   const minAmount0 = amount0InWei
     .mul(bn(100 - Math.floor((params.slippage || DEFAULT_SLIPPAGE) * 100)))

src/mira/swap.ts

@@ -90,7 +90,7 @@ export const swapExactInput = async (
   if (!amountOutWei) {
     throw new Error('Failed to calculate output amount');
   }
-  
+
   const minAmountOut = amountOutWei
     .mul(bn(100 - Math.floor((params.slippage || DEFAULT_SLIPPAGE) * 100)))
     .div(bn(100));

src/utils/models.ts

@@ -0,0 +1,8 @@
+export const modelMapping = {
+  'gpt-4o': 'openai',
+  'gpt-4o-mini': 'openai',
+  'gemini-1.5-flash': 'gemini',
+  'gemini-2.0-flash-exp': 'gemini',
+  'claude-3-5-sonnet-latest': 'anthropic',
+  'claude-3-5-haiku-latest': 'anthropic',
+};

test/add_liquidity.test.ts

@@ -7,26 +7,20 @@ beforeEach(() => {
   agent = createTestAgent();
 });
 
-test(
-  'add liquidity',
-  async () => {
-    console.log(
-      await agent.addLiquidity({
-        amount0: '0.0001',
-        asset0Symbol: 'ETH',
-        asset1Symbol: 'USDT',
-      }),
-    );
-  },
-);
+test('add liquidity', async () => {
+  console.log(
+    await agent.addLiquidity({
+      amount0: '0.0001',
+      asset0Symbol: 'ETH',
+      asset1Symbol: 'USDT',
+    }),
+  );
+});
 
-test(
-  'add liquidity via natural language',
-  async () => {
-    console.log(
-      await agent.execute(
-        'Add liquidity into USDC/USDT pool for 0.1 USDC with 5% slippage',
-      ),
-    );
-  },
-);
+test('add liquidity via natural language', async () => {
+  console.log(
+    await agent.execute(
+      'Add liquidity into USDC/USDT pool for 0.1 USDC with 5% slippage',
+    ),
+  );
+});