updated flash privilige settings and improved token support

src/langchain/index.ts

@@ -9,6 +9,7 @@ import {
   SolanaAgentKit,
 } from "../index";
 import { create_image, FEE_TIERS, generateOrdersfromPattern } from "../tools";
+import { marketTokenMap } from "../utils/flashUtils";
 
 export class SolanaBalanceTool extends Tool {
   name = "solana_balance";
@@ -798,10 +799,14 @@ export class SolanaFlashOpenTrade extends Tool {
       if (!parsedInput.token) {
         throw new Error("Token is required, received: " + parsedInput.token);
       }
-      if (!["SOL", "BTC", "ETH", "USDC"].includes(parsedInput.token)) {
+      if (!Object.keys(marketTokenMap).includes(parsedInput.token)) {
         throw new Error(
-          'Token must be one of ["SOL", "BTC", "ETH", "USDC"], received: ' +
-            parsedInput.token,
+          "Token must be one of " +
+            Object.keys(marketTokenMap).join(", ") +
+            ", received: " +
+            parsedInput.token +
+            "\n" +
+            "Please check https://beast.flash.trade/ for the list of supported tokens",
         );
       }
       if (!["long", "short"].includes(parsedInput.type)) {

src/tools/flash_close_trade.ts

@@ -1,5 +1,5 @@
 import { ComputeBudgetProgram } from "@solana/web3.js";
-import { PoolConfig, Privilege, Side } from "flash-sdk";
+import { PoolConfig, Side } from "flash-sdk";
 import { BN } from "@coral-xyz/anchor";
 import { SolanaAgentKit } from "../index";
 import {
@@ -9,6 +9,7 @@ import {
   getNftTradingAccountInfo,
   fetchOraclePrice,
   createPerpClient,
+  get_flash_privilege,
 } from "../utils/flashUtils";
 import { FlashCloseTradeParams } from "../types";
 
@@ -93,7 +94,7 @@ export async function flashCloseTrade(
       priceWithSlippage,
       sideEnum,
       poolConfig,
-      Privilege.Referral,
+      get_flash_privilege(),
       tradingAccounts.nftTradingAccountPk,
       tradingAccounts.nftReferralAccountPK,
       tradingAccounts.nftOwnerRebateTokenAccountPk,

src/tools/flash_open_trade.ts

@@ -3,7 +3,6 @@ import {
   PerpetualsClient,
   OraclePrice,
   PoolConfig,
-  Privilege,
   Side,
   CustodyAccount,
   Custody,
@@ -18,6 +17,7 @@ import {
   OPEN_POSITION_CU,
   fetchOraclePrice,
   createPerpClient,
+  get_flash_privilege,
 } from "../utils/flashUtils";
 import { FlashTradeParams } from "../types";
 
@@ -141,7 +141,7 @@ export async function flashOpenTrade(
       positionSize,
       side === "long" ? Side.Long : Side.Short,
       poolConfig,
-      Privilege.Referral,
+      get_flash_privilege(),
       tradingAccounts.nftTradingAccountPk,
       tradingAccounts.nftReferralAccountPK,
       tradingAccounts.nftOwnerRebateTokenAccountPk!,

src/utils/flashUtils.ts

@@ -1,7 +1,13 @@
 import { HermesClient } from "@pythnetwork/hermes-client";
 import { OraclePrice } from "flash-sdk";
 import { AnchorProvider, BN, Wallet } from "@coral-xyz/anchor";
-import { PoolConfig, Token, Referral, PerpetualsClient } from "flash-sdk";
+import {
+  PoolConfig,
+  Token,
+  Referral,
+  PerpetualsClient,
+  Privilege,
+} from "flash-sdk";
 import { Cluster, PublicKey, Connection, Keypair } from "@solana/web3.js";
 import { getAssociatedTokenAddressSync } from "@solana/spl-token";
 
@@ -278,3 +284,16 @@ export function createPerpClient(
     {},
   );
 }
+
+const FLASH_PRIVILEGE = process.env.FLASH_PRIVILEGE || "None";
+
+export function get_flash_privilege(): Privilege {
+  switch (FLASH_PRIVILEGE.toLowerCase()) {
+    case "referral":
+      return Privilege.Referral;
+    case "nft":
+      return Privilege.NFT;
+    default:
+      return Privilege.None;
+  }
+}