From 97e96730896551e64f69b16543816c219d3c17fe Mon Sep 17 00:00:00 2001 From: michaelessiet Date: Thu, 16 Jan 2025 19:39:47 +0100 Subject: [PATCH] feat: add more drift actions --- src/actions/drift/availableMarkets.ts | 55 ++++ src/actions/drift/createDriftUserAccount.ts | 7 +- src/actions/drift/createVault.ts | 9 +- src/actions/drift/depositIntoVault.ts | 4 +- .../drift/depositToDriftUserAccount.ts | 2 +- .../requestUnstakeFromDriftInsuranceFund.ts | 62 ++++ .../drift/requestWithdrawalFromVault.ts | 4 +- .../drift/stakeToDriftInsuranceFund.ts | 59 ++++ src/actions/drift/swapSpotToken.ts | 76 +++++ src/actions/drift/tradeDelegatedDriftVault.ts | 13 +- src/actions/drift/tradePerpAccount.ts | 21 +- .../drift/unstakeFromDriftInsuranceFund.ts | 51 ++++ src/actions/drift/updateDriftVaultDelegate.ts | 4 +- src/actions/drift/updateVault.ts | 31 +- src/actions/drift/withdrawFromDriftAccount.ts | 2 +- src/actions/drift/withdrawFromVault.ts | 2 +- src/actions/index.ts | 11 + src/agent/index.ts | 50 ++++ .../request_unstake_from_insurance_fund.ts | 37 +++ .../drift/stake_to_insurance_fund.ts | 37 +++ src/langchain/drift/swap_spot_token.ts | 37 +++ .../drift/unstake_from_insurance_fund.ts | 32 +++ src/tools/drift/drift.ts | 270 +++++++++++++++++- src/tools/drift/drift_vault.ts | 30 +- 24 files changed, 863 insertions(+), 43 deletions(-) create mode 100644 src/actions/drift/availableMarkets.ts create mode 100644 src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts create mode 100644 src/actions/drift/stakeToDriftInsuranceFund.ts create mode 100644 src/actions/drift/swapSpotToken.ts create mode 100644 src/actions/drift/unstakeFromDriftInsuranceFund.ts create mode 100644 src/langchain/drift/request_unstake_from_insurance_fund.ts create mode 100644 src/langchain/drift/stake_to_insurance_fund.ts create mode 100644 src/langchain/drift/swap_spot_token.ts create mode 100644 src/langchain/drift/unstake_from_insurance_fund.ts diff --git a/src/actions/drift/availableMarkets.ts b/src/actions/drift/availableMarkets.ts new file mode 100644 index 0000000..85bd174 --- /dev/null +++ b/src/actions/drift/availableMarkets.ts @@ -0,0 +1,55 @@ +import { MainnetSpotMarkets } from "@drift-labs/sdk"; +import type { Action } from "../../types"; +import { z } from "zod"; +import { + getAvailableDriftPerpMarkets, + getAvailableDriftSpotMarkets, +} from "../../tools"; + +const availableDriftMarketsAction: Action = { + name: "AVAILABLE_DRIFT_MARKETS", + description: "Get a list of available drift markets", + similes: [ + "get drift markets", + "drift markets", + "available drift markets", + "get available drift perp markets", + "get available spot markets on drift", + ], + examples: [ + [ + { + input: { + marketType: "spot", + }, + output: { + status: "success", + message: `The list of available spot markets are ${MainnetSpotMarkets.map((v) => v.symbol).join(", ")}`, + data: MainnetSpotMarkets, + }, + explanation: "Get the list of available spot markets/tokens on drift", + }, + ], + ], + schema: z.object({ + marketType: z + .enum(["spot", "perp"]) + .describe("Type of market to get") + .optional(), + }), + handler: async (agent, input) => { + switch (input.marketType) { + case "perp": + return getAvailableDriftPerpMarkets(); + case "spot": + return getAvailableDriftSpotMarkets(); + default: + return { + spot: getAvailableDriftSpotMarkets(), + perp: getAvailableDriftPerpMarkets(), + }; + } + }, +}; + +export default availableDriftMarketsAction; diff --git a/src/actions/drift/createDriftUserAccount.ts b/src/actions/drift/createDriftUserAccount.ts index 32c62ef..43638c3 100644 --- a/src/actions/drift/createDriftUserAccount.ts +++ b/src/actions/drift/createDriftUserAccount.ts @@ -27,7 +27,12 @@ const createDriftUserAccountAction: Action = { ], ], schema: z.object({ - amount: z.number().positive().describe("Amount of the token to deposit"), + amount: z + .number() + .positive() + .describe( + "Amount of the token to deposit. In normal token amounts e.g 50 SOL, 100 USDC, etc", + ), symbol: z.string().describe("Symbol of the token to deposit"), }), handler: async (agent, input) => { diff --git a/src/actions/drift/createVault.ts b/src/actions/drift/createVault.ts index 26ac59b..e09e07d 100644 --- a/src/actions/drift/createVault.ts +++ b/src/actions/drift/createVault.ts @@ -53,9 +53,14 @@ const createDriftVaultAction: Action = { .int() .min(100, "Max tokens must be at least 100") .describe( - "The maximum amount of tokens the vault will be accomodating. For example some vaults have a cap at 10 million USDC", + "The maximum amount of tokens the vault will be accomodating. For example some vaults have a cap at 10 million USDC. This amount should be normal token amounts e.g 50 SOL, 100 USDC, etc", + ), + minDepositAmount: z + .number() + .positive() + .describe( + "Minimum deposit amount in normal token values e.g 50 SOL, 100 USDC, etc", ), - minDepositAmount: z.number().positive().describe("Minimum deposit amount"), managementFee: z .number() .positive() diff --git a/src/actions/drift/depositIntoVault.ts b/src/actions/drift/depositIntoVault.ts index eddb9f0..ca14c53 100644 --- a/src/actions/drift/depositIntoVault.ts +++ b/src/actions/drift/depositIntoVault.ts @@ -28,7 +28,9 @@ const depositIntoDriftVaultAction: Action = { amount: z .number() .positive() - .describe("The amount in tokens you'd like to deposit into the vault"), + .describe( + "The amount in tokens you'd like to deposit into the vault in normal token amounts e.g 50 SOL, 100 USDC, etc", + ), }), handler: async (agent, input) => { try { diff --git a/src/actions/drift/depositToDriftUserAccount.ts b/src/actions/drift/depositToDriftUserAccount.ts index af152c3..3e0dc7b 100644 --- a/src/actions/drift/depositToDriftUserAccount.ts +++ b/src/actions/drift/depositToDriftUserAccount.ts @@ -34,7 +34,7 @@ const depositToDriftUserAccountAction: Action = { .number() .positive() .describe( - "The amount in tokens you'd like to deposit into your drift user account", + "The amount in tokens you'd like to deposit into your drift user account in normal token amounts e.g 50 SOL, 100 USDC, etc", ), symbol: z .string() diff --git a/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts b/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts new file mode 100644 index 0000000..fc1b645 --- /dev/null +++ b/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts @@ -0,0 +1,62 @@ +import { z } from "zod"; +import type { Action } from "../../types"; +import { requestUnstakeFromDriftInsuranceFund } from "../../tools"; + +const requestUnstakeFromDriftInsuranceFundAction: Action = { + name: "REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION", + description: + "Request to unstake a certain amount of a token from the Drift Insurance Fund", + similes: [ + "request an unstake from the drift insurance fund", + "unstake an amount from the drift insurance fund", + "ask to unstake a certain amount from the drift insurance fund", + ], + examples: [ + [ + { + input: { + amount: 100, + symbol: "SOL", + }, + output: { + status: "success", + message: "Requested to unstake 100 SOL from the Drift Insurance Fund", + signature: "4FdasklhiIHyOI", + }, + explanation: "Request to unstake 100 SOL from the Drift Insurance Fund", + }, + ], + ], + schema: z.object({ + amount: z + .number() + .positive() + .describe("Amount to unstake in normal units e.g 50 === 50 SOL"), + symbol: z.string().describe("Symbol of the token to unstake"), + }), + handler: async (agent, input) => { + try { + const tx = await requestUnstakeFromDriftInsuranceFund( + agent, + input.amount, + input.symbol, + ); + + return { + status: "success", + message: `Requested to unstake ${input.amount} ${input.symbol} from the Drift Insurance Fund`, + data: { + signature: tx, + }, + }; + } catch (e) { + return { + status: "error", + // @ts-expect-error error is not a string + message: e.message, + }; + } + }, +}; + +export default requestUnstakeFromDriftInsuranceFundAction; diff --git a/src/actions/drift/requestWithdrawalFromVault.ts b/src/actions/drift/requestWithdrawalFromVault.ts index df6939c..9dd1641 100644 --- a/src/actions/drift/requestWithdrawalFromVault.ts +++ b/src/actions/drift/requestWithdrawalFromVault.ts @@ -29,7 +29,9 @@ const requestWithdrawalFromVaultAction: Action = { amount: z .number() .positive() - .describe("Amount of shares you would like to withdraw from the vault"), + .describe( + "Amount of shares you would like to withdraw from the vault in normal token amounts e.g 50 SOL, 100 USDC, etc", + ), }), handler: async (agent: SolanaAgentKit, input) => { try { diff --git a/src/actions/drift/stakeToDriftInsuranceFund.ts b/src/actions/drift/stakeToDriftInsuranceFund.ts new file mode 100644 index 0000000..39cf179 --- /dev/null +++ b/src/actions/drift/stakeToDriftInsuranceFund.ts @@ -0,0 +1,59 @@ +import { z } from "zod"; +import type { Action } from "../../types"; +import { stakeToDriftInsuranceFund } from "../../tools"; + +const stakeToDriftInsuranceFundAction: Action = { + name: "STAKE_TO_DRIFT_INSURANCE_FUND_ACTION", + description: "Stake a token to Drift Insurance Fund", + similes: ["Stake a token to Drift Insurance Fund"], + examples: [ + [ + { + input: { + amount: 100, + symbol: "SOL", + }, + output: { + status: "success", + message: "Staked 100 SOL to the Drift Insurance Fund", + data: { + signature: "signature", + }, + }, + explanation: "Stake 100 SOL to the Drift Insurance Fund", + }, + ], + ], + schema: z.object({ + amount: z + .number() + .positive() + .describe("Amount to stake in normal units e.g 50 === 50 SOL"), + symbol: z.string().describe("Symbol of the token stake"), + }), + handler: async (agent, input) => { + try { + const tx = await stakeToDriftInsuranceFund( + agent, + input.amount, + input.symbol, + ); + + return { + status: "sucess", + message: `Staked ${input.amount} ${input.symbol} to the Drift Insurance Fund`, + data: { + signature: tx, + }, + }; + } catch (error) { + return { + status: "error", + // @ts-expect-error error is not a string + message: error.message, + }; + } + }, +}; + +export default stakeToDriftInsuranceFundAction; diff --git a/src/actions/drift/swapSpotToken.ts b/src/actions/drift/swapSpotToken.ts new file mode 100644 index 0000000..a076cb8 --- /dev/null +++ b/src/actions/drift/swapSpotToken.ts @@ -0,0 +1,76 @@ +import { z } from "zod"; +import type { Action } from "../../types"; +import { swapSpotToken } from "../../tools"; + +const driftSpotTokenSwapAction: Action = { + name: "DRIFT_SPOT_TOKEN_SWAP_ACTION", + description: "Swap a spot token for another spot token on Drift", + similes: [ + "swap a token for another token on drift", + "exchange a token for another token on drift", + "trade a token for another token on drift", + ], + examples: [ + [ + { + input: { + fromSymbol: "SOL", + toSymbol: "USDC", + fromAmount: 100, + }, + output: { + status: "success", + message: "Swapped 100 SOL for USDC on Drift", + signature: "4FdasklhiIHyOI", + }, + explanation: "Swap 100 SOL for USDC on Drift", + }, + ], + ], + schema: z.object({ + fromSymbol: z.string().describe("Symbol of the token to swap from"), + toSymbol: z.string().describe("Symbol of the token to swap to"), + fromAmount: z + .number() + .positive() + .describe("Amount to swap from in normal units e.g 50 === 50 SOL") + .optional(), + toAmount: z + .number() + .positive() + .describe("Amount to swap to in normal units e.g 5000 === 5000 USDC") + .optional(), + slippage: z + .number() + .positive() + .describe("Slippage tolerance in percentage e.g 0.5 === 0.5%") + .optional(), + }), + handler: async (agent, input) => { + try { + const tx = await swapSpotToken(agent, { + fromSymbol: input.fromSymbol, + toSymbol: input.toSymbol, + fromAmount: input.fromAmount, + toAmount: input.toAmount, + slippage: input.slippage, + }); + + return { + status: "success", + message: `Swapped ${input.fromAmount} ${input.fromSymbol} for ${input.toAmount} ${input.toSymbol} on Drift`, + data: { + signature: tx, + }, + }; + } catch (e) { + return { + status: "error", + // @ts-expect-error error is not a string + message: e.message, + }; + } + }, +}; + +export default driftSpotTokenSwapAction; diff --git a/src/actions/drift/tradeDelegatedDriftVault.ts b/src/actions/drift/tradeDelegatedDriftVault.ts index c85d8b3..5d28a84 100644 --- a/src/actions/drift/tradeDelegatedDriftVault.ts +++ b/src/actions/drift/tradeDelegatedDriftVault.ts @@ -64,11 +64,20 @@ const tradeDelegatedDriftVaultAction: Action = { ], schema: z.object({ vaultAddress: z.string().describe("Address of the Drift vault to trade in"), - amount: z.number().positive().describe("Amount to trade"), + amount: z + .number() + .positive() + .describe( + "Amount to trade in normal token amounts e.g 50 SOL, 100 USDC, etc", + ), symbol: z.string().describe("Symbol of the token to trade"), action: z.enum(["long", "short"]).describe("Trade action - long or short"), type: z.enum(["market", "limit"]).describe("Trade type - market or limit"), - price: z.number().positive().optional().describe("Price for limit order"), + price: z + .number() + .positive() + .optional() + .describe("USD price for limit order"), }), handler: async (agent: SolanaAgentKit, input) => { try { diff --git a/src/actions/drift/tradePerpAccount.ts b/src/actions/drift/tradePerpAccount.ts index ae1db7e..15a53b4 100644 --- a/src/actions/drift/tradePerpAccount.ts +++ b/src/actions/drift/tradePerpAccount.ts @@ -46,14 +46,27 @@ export const tradeDriftPerpAccountAction: Action = { ], ], schema: z.object({ - amount: z.number().positive(), + amount: z + .number() + .positive() + .describe( + "The amount of the token to trade in normal token amounts e.g 50 SOL, 100 USDC", + ), symbol: z .string() .toUpperCase() .describe("Symbol of the token to open a position on "), - action: z.enum(["long", "short"]), - type: z.enum(["market", "limit"]), - price: z.number().positive().optional(), + action: z + .enum(["long", "short"]) + .describe( + "The action you would like to carry out whether it be a long or a short", + ), + type: z + .enum(["market", "limit"]) + .describe( + "The type of trade you would like to open, market or limit order", + ), + price: z.number().positive().optional().describe("USD price of the token"), }), handler: async (agent, input) => { try { diff --git a/src/actions/drift/unstakeFromDriftInsuranceFund.ts b/src/actions/drift/unstakeFromDriftInsuranceFund.ts new file mode 100644 index 0000000..4d9516e --- /dev/null +++ b/src/actions/drift/unstakeFromDriftInsuranceFund.ts @@ -0,0 +1,51 @@ +import { z } from "zod"; +import type { Action } from "../../types"; +import { unstakeFromDriftInsuranceFund } from "../../tools"; + +const unstakeFromDriftInsuranceFundAction: Action = { + name: "UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION", + description: + "Unstake requested unstake amount from the Drift Insurance fund once the cool period has elapsed", + similes: [ + "unstake from the drift insurance fund", + "withdraw from the drift insurance fund", + "take out funds from the drift insurance fund", + ], + examples: [ + [ + { + input: { + symbol: "SOL", + }, + output: { + status: "success", + message: "Unstaked your SOL from the Drift Insurance Fund", + signature: "4FdasklhiIHyOI", + }, + explanation: "Unstake SOL from the Drift Insurance Fund", + }, + ], + ], + schema: z.object({ + symbol: z.string().describe("Symbol of the token to unstake"), + }), + handler: async (agent, input) => { + try { + const tx = await unstakeFromDriftInsuranceFund(agent, input.symbol); + + return { + status: "success", + message: `Unstaked your ${input.symbol} from the Drift Insurance Fund`, + signature: tx, + }; + } catch (e) { + return { + status: "error", + // @ts-expect-error error is not a string + message: e.message, + }; + } + }, +}; + +export default unstakeFromDriftInsuranceFundAction; diff --git a/src/actions/drift/updateDriftVaultDelegate.ts b/src/actions/drift/updateDriftVaultDelegate.ts index defc7a6..d203e0d 100644 --- a/src/actions/drift/updateDriftVaultDelegate.ts +++ b/src/actions/drift/updateDriftVaultDelegate.ts @@ -24,8 +24,8 @@ const updateDriftVaultDelegateAction: Action = { ], ], schema: z.object({ - vaultAddress: z.string(), - newDelegate: z.string(), + vaultAddress: z.string().describe("vault's address"), + newDelegate: z.string().describe("new address to delegate the vault to"), }), handler: async (agent, input) => { try { diff --git a/src/actions/drift/updateVault.ts b/src/actions/drift/updateVault.ts index 4d0f66e..2d0abed 100644 --- a/src/actions/drift/updateVault.ts +++ b/src/actions/drift/updateVault.ts @@ -41,16 +41,35 @@ const updateDriftVaultAction: Action = { redeemPeriod: z .number() .int() - .min(1, "Redeem period must be at least 1") + .min(1, "Redeem period must be at least 1 day") .optional(), maxTokens: z .number() .int() - .min(100, "Max tokens must be at least 100") - .optional(), - minDepositAmount: z.number().positive().optional(), - managementFee: z.number().positive().max(20).optional(), - profitShare: z.number().positive().max(90).optional(), + .min(100, "Max tokens must be at least be 100 units") + .optional() + .describe( + "The maximum number of tokens the vault is willing to accept and manage", + ), + minDepositAmount: z + .number() + .positive() + .optional() + .describe( + "The minimum amount that is allowed to be deposited into the vault in normal token amounts e.g 10 USDC", + ), + managementFee: z + .number() + .positive() + .max(20) + .optional() + .describe("The percentage fee the vault takes for asset management"), + profitShare: z + .number() + .positive() + .max(90) + .optional() + .describe("Profit share in percentage e.g 2 === 2%"), handleRate: z.number().optional(), permissioned: z .boolean() diff --git a/src/actions/drift/withdrawFromDriftAccount.ts b/src/actions/drift/withdrawFromDriftAccount.ts index 00d72b6..693a294 100644 --- a/src/actions/drift/withdrawFromDriftAccount.ts +++ b/src/actions/drift/withdrawFromDriftAccount.ts @@ -36,7 +36,7 @@ const withdrawFromDriftAccountAction: Action = { .number() .positive() .describe( - "The amount in tokens you'd like to withdraw from your drift account", + "The amount in tokens you'd like to withdraw from your drift account in normal token amounts, e.g 50 SOL, 100 USDC, etc", ), symbol: z .string() diff --git a/src/actions/drift/withdrawFromVault.ts b/src/actions/drift/withdrawFromVault.ts index b6007f2..4b4318b 100644 --- a/src/actions/drift/withdrawFromVault.ts +++ b/src/actions/drift/withdrawFromVault.ts @@ -25,7 +25,7 @@ const withdrawFromVaultAction: Action = { ], ], schema: z.object({ - vaultAddress: z.string(), + vaultAddress: z.string().describe("Vault's address"), }), handler: async (agent: SolanaAgentKit, input) => { try { diff --git a/src/actions/index.ts b/src/actions/index.ts index 29cf233..5b4324e 100644 --- a/src/actions/index.ts +++ b/src/actions/index.ts @@ -59,6 +59,11 @@ import withdrawFromDriftAccountAction from "./drift/withdrawFromDriftAccount"; import driftUserAccountInfoAction from "./drift/driftUserAccountInfo"; import deriveDriftVaultAddressAction from "./drift/deriveVaultAddress"; import updateDriftVaultDelegateAction from "./drift/updateDriftVaultDelegate"; +import availableDriftMarketsAction from "./drift/availableMarkets"; +import stakeToDriftInsuranceFundAction from "./drift/stakeToDriftInsuranceFund"; +import requestUnstakeFromDriftInsuranceFundAction from "./drift/requestUnstakeFromDriftInsuranceFund"; +import unstakeFromDriftInsuranceFundAction from "./drift/unstakeFromDriftInsuranceFund"; +import driftSpotTokenSwapAction from "./drift/swapSpotToken"; export const ACTIONS = { WALLET_ADDRESS_ACTION: getWalletAddressAction, @@ -123,6 +128,12 @@ export const ACTIONS = { DRIFT_USER_ACCOUNT_INFO_ACTION: driftUserAccountInfoAction, DERIVE_DRIFT_VAULT_ADDRESS_ACTION: deriveDriftVaultAddressAction, UPDATE_DRIFT_VAULT_DELEGATE_ACTION: updateDriftVaultDelegateAction, + AVAILABLE_DRIFT_MARKETS_ACTION: availableDriftMarketsAction, + STAKE_TO_DRIFT_INSURANCE_FUND_ACTION: stakeToDriftInsuranceFundAction, + REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION: + requestUnstakeFromDriftInsuranceFundAction, + UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION: unstakeFromDriftInsuranceFundAction, + DRIFT_SPOT_TOKEN_SWAP_ACTION: driftSpotTokenSwapAction, }; export type { Action, ActionExample, Handler } from "../types/action"; diff --git a/src/agent/index.ts b/src/agent/index.ts index d087931..0f6c834 100644 --- a/src/agent/index.ts +++ b/src/agent/index.ts @@ -98,6 +98,12 @@ import { withdrawFromDriftVault, updateVaultDelegate, get_token_balance, + getAvailableDriftSpotMarkets, + getAvailableDriftPerpMarkets, + stakeToDriftInsuranceFund, + requestUnstakeFromDriftInsuranceFund, + unstakeFromDriftInsuranceFund, + swapSpotToken, } from "../tools"; import { Config, @@ -821,4 +827,48 @@ export class SolanaAgentKit { async updateDriftVaultDelegate(vaultAddress: string, delegate: string) { return await updateVaultDelegate(this, vaultAddress, delegate); } + getAvailableDriftMarkets(type?: "spot" | "perp") { + switch (type) { + case "spot": + return getAvailableDriftSpotMarkets(); + case "perp": + return getAvailableDriftPerpMarkets(); + default: + return { + spot: getAvailableDriftSpotMarkets(), + perp: getAvailableDriftPerpMarkets(), + }; + } + } + async stakeToDriftInsuranceFund(amount: number, symbol: string) { + return await stakeToDriftInsuranceFund(this, amount, symbol); + } + async requestUnstakeFromDriftInsuranceFund(amount: number, symbol: string) { + return await requestUnstakeFromDriftInsuranceFund(this, amount, symbol); + } + async unstakeFromDriftInsuranceFund(symbol: string) { + return await unstakeFromDriftInsuranceFund(this, symbol); + } + async driftSpotTokenSwap( + params: { + fromSymbol: string; + toSymbol: string; + slippage?: number; + } & ( + | { + toAmount: number; + } + | { fromAmount: number } + ), + ) { + return await swapSpotToken(this, { + fromSymbol: params.fromSymbol, + toSymbol: params.toSymbol, + // @ts-expect-error - fromAmount and toAmount are mutually exclusive + fromAmount: params.fromAmount, + // @ts-expect-error - fromAmount and toAmount are mutually exclusive + toAmount: params.toAmount, + slippage: params.slippage, + }); + } } diff --git a/src/langchain/drift/request_unstake_from_insurance_fund.ts b/src/langchain/drift/request_unstake_from_insurance_fund.ts new file mode 100644 index 0000000..edeb630 --- /dev/null +++ b/src/langchain/drift/request_unstake_from_insurance_fund.ts @@ -0,0 +1,37 @@ +import { Tool } from "langchain/tools"; +import type { SolanaAgentKit } from "../../agent"; + +export class SolanaRequestUnstakeFromDriftInsuranceFundTool extends Tool { + name = "request_unstake_from_drift_insurance_fund"; + description = `Request to unstake tokens from Drift Insurance Fund. + + Inputs (JSON string): + - amount: number, amount to unstake (required) + - symbol: string, token symbol (required)`; + + constructor(private solanaKit: SolanaAgentKit) { + super(); + } + + protected async _call(input: string): Promise { + try { + const parsedInput = JSON.parse(input); + const tx = await this.solanaKit.requestUnstakeFromDriftInsuranceFund( + parsedInput.amount, + parsedInput.symbol, + ); + + return JSON.stringify({ + status: "success", + message: `Requested unstake of ${parsedInput.amount} ${parsedInput.symbol} from the Drift Insurance Fund`, + signature: tx, + }); + } catch (error: any) { + return JSON.stringify({ + status: "error", + message: error.message, + code: error.code || "REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ERROR", + }); + } + } +} diff --git a/src/langchain/drift/stake_to_insurance_fund.ts b/src/langchain/drift/stake_to_insurance_fund.ts new file mode 100644 index 0000000..09cc363 --- /dev/null +++ b/src/langchain/drift/stake_to_insurance_fund.ts @@ -0,0 +1,37 @@ +import { Tool } from "langchain/tools"; +import type { SolanaAgentKit } from "../../agent"; + +export class SolanaStakeToDriftInsuranceFundTool extends Tool { + name = "stake_to_drift_insurance_fund"; + description = `Stake a token to Drift Insurance Fund. + + Inputs (JSON string): + - amount: number, amount to stake (required) + - symbol: string, token symbol (required)`; + + constructor(private solanaKit: SolanaAgentKit) { + super(); + } + + protected async _call(input: string): Promise { + try { + const parsedInput = JSON.parse(input); + const tx = await this.solanaKit.stakeToDriftInsuranceFund( + parsedInput.amount, + parsedInput.symbol, + ); + + return JSON.stringify({ + status: "success", + message: `Staked ${parsedInput.amount} ${parsedInput.symbol} to the Drift Insurance Fund`, + signature: tx, + }); + } catch (error: any) { + return JSON.stringify({ + status: "error", + message: error.message, + code: error.code || "STAKE_TO_DRIFT_INSURANCE_FUND_ERROR", + }); + } + } +} diff --git a/src/langchain/drift/swap_spot_token.ts b/src/langchain/drift/swap_spot_token.ts new file mode 100644 index 0000000..f5f7314 --- /dev/null +++ b/src/langchain/drift/swap_spot_token.ts @@ -0,0 +1,37 @@ +import { Tool } from "langchain/tools"; +import type { SolanaAgentKit } from "../../agent"; + +export class SolanaDriftSpotTokenSwapTool extends Tool { + name = "drift_spot_token_swap"; + description = `Swap spot tokens on Drift protocol. + + Inputs (JSON string): + - fromSymbol: string, symbol of token to swap from (required) + - toSymbol: string, symbol of token to swap to (required) + - fromAmount: number, amount to swap from (optional) required if toAmount is not provided + - toAmount: number, amount to swap to (optional) required if fromAmount is not provided + - slippage: number, slippage tolerance in percentage (optional)`; + + constructor(private solanaKit: SolanaAgentKit) { + super(); + } + + protected async _call(input: string): Promise { + try { + const parsedInput = JSON.parse(input); + const tx = await this.solanaKit.driftSpotTokenSwap(parsedInput); + + return JSON.stringify({ + status: "success", + message: `Swapped ${parsedInput.fromAmount} ${parsedInput.fromSymbol} for ${parsedInput.toSymbol}`, + signature: tx, + }); + } catch (error: any) { + return JSON.stringify({ + status: "error", + message: error.message, + code: error.code || "DRIFT_SPOT_TOKEN_SWAP_ERROR", + }); + } + } +} diff --git a/src/langchain/drift/unstake_from_insurance_fund.ts b/src/langchain/drift/unstake_from_insurance_fund.ts new file mode 100644 index 0000000..b610aae --- /dev/null +++ b/src/langchain/drift/unstake_from_insurance_fund.ts @@ -0,0 +1,32 @@ +import { Tool } from "langchain/tools"; +import type { SolanaAgentKit } from "../../agent"; + +export class SolanaUnstakeFromDriftInsuranceFundTool extends Tool { + name = "unstake_from_drift_insurance_fund"; + description = `Unstake tokens from Drift Insurance Fund after request period has elapsed. + + Inputs (JSON string): + - symbol: string, token symbol (required)`; + + constructor(private solanaKit: SolanaAgentKit) { + super(); + } + + protected async _call(input: string): Promise { + try { + const tx = await this.solanaKit.unstakeFromDriftInsuranceFund(input); + + return JSON.stringify({ + status: "success", + message: `Unstaked ${input} from the Drift Insurance Fund`, + signature: tx, + }); + } catch (error: any) { + return JSON.stringify({ + status: "error", + message: error.message, + code: error.code || "UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ERROR", + }); + } + } +} diff --git a/src/tools/drift/drift.ts b/src/tools/drift/drift.ts index 1a8b4a9..57999f5 100644 --- a/src/tools/drift/drift.ts +++ b/src/tools/drift/drift.ts @@ -7,6 +7,8 @@ import { getLimitOrderParams, getMarketOrderParams, getUserAccountPublicKeySync, + JupiterClient, + MainnetPerpMarkets, MainnetSpotMarkets, numberToSafeBN, PositionDirection, @@ -115,7 +117,10 @@ export async function createDriftUserAccount( ); if (!token) { - throw new Error(`Token with symbol ${symbol} not found`); + throw new Error(`Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")} + `); } if (!userAccountExists) { @@ -171,7 +176,11 @@ export async function depositToDriftUserAccount( ); if (!token) { - throw new Error(`Token with symbol ${symbol} not found`); + throw new Error( + `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); } if (!userAccountExists) { @@ -237,7 +246,11 @@ export async function withdrawFromDriftUserAccount( ); if (!token) { - throw new Error(`Token with symbol ${symbol} not found`); + throw new Error( + `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); } const withdrawAmount = numberToSafeBN(amount, token.precision); @@ -313,7 +326,11 @@ export async function driftPerpTrade( ); if (!market) { - throw new Error(`Token with symbol ${params.symbol} not found`); + throw new Error( + `Token with symbol ${params.symbol} not found. Here's a list of available perp markets: ${MainnetPerpMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); } const baseAssetPrice = driftClient.getOracleDataForPerpMarket( @@ -435,10 +452,9 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) { })); const spotPositions = account.spotPositions.map((pos) => ({ ...pos, - scaledBalance: convertToNumber(pos.scaledBalance, BASE_PRECISION), - cumulativeDeposits: convertToNumber( - pos.cumulativeDeposits, - BASE_PRECISION, + availableBalance: convertToNumber( + pos.scaledBalance, + MainnetSpotMarkets[pos.marketIndex].precision, ), symbol: MainnetSpotMarkets.find((v) => v.marketIndex === pos.marketIndex) ?.symbol, @@ -448,8 +464,6 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) { ...account, name: account.name, authority: account.authority, - totalDeposits: `$${convertToNumber(account.totalDeposits, QUOTE_PRECISION)}`, - totalWithdraws: `$${convertToNumber(account.totalWithdraws, QUOTE_PRECISION)}`, settledPerpPnl: `$${convertToNumber(account.settledPerpPnl, QUOTE_PRECISION)}`, lastActiveSlot: account.lastActiveSlot.toNumber(), perpPositions, @@ -460,3 +474,239 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) { throw new Error(`Failed to check user account: ${e.message}`); } } + +/** + * Get available spot markets on drift protocol + */ +export function getAvailableDriftSpotMarkets() { + return MainnetSpotMarkets; +} + +/** + * Get available perp markets on drift protocol + */ +export function getAvailableDriftPerpMarkets() { + return MainnetPerpMarkets; +} + +/** + * Stake a token to the drift insurance fund + * @param agent + * @param amount + * @param symbol + */ +export async function stakeToDriftInsuranceFund( + agent: SolanaAgentKit, + amount: number, + symbol: string, +) { + try { + const { cleanUp, driftClient } = await initClients(agent); + const token = MainnetSpotMarkets.find( + (v) => v.symbol === symbol.toUpperCase(), + ); + + if (!token) { + throw new Error( + `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); + } + + const signature = await driftClient.addInsuranceFundStake({ + amount: numberToSafeBN(amount, token.precision), + marketIndex: token.marketIndex, + collateralAccountPublicKey: getAssociatedTokenAddressSync( + token.mint, + agent.wallet.publicKey, + ), + txParams: { + computeUnitsPrice: 0.000002 * 1000000 * 1000000, + }, + }); + + await cleanUp(); + return signature; + } catch (e) { + // @ts-expect-error - error message is a string + throw new Error(`Failed to get APYs: ${e.message}`); + } +} + +/** + * Request an unstake from the drift insurance fund + * @param agent + * @param amount + * @param symbol + */ +export async function requestUnstakeFromDriftInsuranceFund( + agent: SolanaAgentKit, + amount: number, + symbol: string, +) { + try { + const { driftClient, cleanUp } = await initClients(agent); + const token = MainnetSpotMarkets.find( + (v) => v.symbol === symbol.toUpperCase(), + ); + + if (!token) { + throw new Error( + `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); + } + + const signature = await driftClient.requestRemoveInsuranceFundStake( + token.marketIndex, + numberToSafeBN(amount, token.precision), + { computeUnitsPrice: 0.000002 * 1000000 * 1000000 }, + ); + + await cleanUp(); + return signature; + } catch (e) { + // @ts-expect-error error message is a string + throw new Error(`Failed to unstake from insurance fund: ${e.message}`); + } +} + +/** + * Unstake requested funds from the drift insurance fund once cool down period is elapsed + * @param agent + * @param symbol + */ +export async function unstakeFromDriftInsuranceFund( + agent: SolanaAgentKit, + symbol: string, +) { + try { + const { driftClient, cleanUp } = await initClients(agent); + const token = MainnetSpotMarkets.find( + (v) => v.symbol === symbol.toUpperCase(), + ); + + if (!token) { + throw new Error( + `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); + } + + const signature = await driftClient.removeInsuranceFundStake( + token.marketIndex, + getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey), + { + computeUnitsPrice: 0.000002 * 1000000 * 1000000, + }, + ); + + await cleanUp(); + return signature; + } catch (e) { + // @ts-expect-error error message is a string + throw new Error(`Failed to unstake from insurance fund: ${e.message}`); + } +} + +/** + * Swap a spot token for another on drift + * @param agent + * @param params + * @param params.fromSymbol symbol of the token to deposit + * @param params.toSymbol symbol of the token to receive + * @param params.fromAmount amount of the token to deposit + * @param params.toAmount amount of the token to receive + */ +export async function swapSpotToken( + agent: SolanaAgentKit, + params: { + fromSymbol: string; + toSymbol: string; + slippage?: number | undefined; + } & ( + | { + fromAmount: number; + } + | { + toAmount: number; + } + ), +) { + try { + const { driftClient, cleanUp } = await initClients(agent); + const fromToken = MainnetSpotMarkets.find( + (v) => v.symbol === params.fromSymbol.toUpperCase(), + ); + const toToken = MainnetSpotMarkets.find( + (v) => v.symbol === params.toSymbol.toUpperCase(), + ); + + if (!fromToken) { + throw new Error( + `Token with symbol ${params.fromSymbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); + } + + if (!toToken) { + throw new Error( + `Token with symbol ${params.toSymbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map( + (v) => v.symbol, + ).join(", ")}`, + ); + } + + let txSig: string; + + // @ts-expect-error - false undefined type conflict + if (params.fromAmount) { + const jupiterClient = new JupiterClient({ connection: agent.connection }); + // @ts-expect-error - false undefined type conflict + const fromAmount = numberToSafeBN(params.fromAmount, fromToken.precision); + const signature = await driftClient.swap({ + amount: fromAmount, + inMarketIndex: fromToken.marketIndex, + outMarketIndex: toToken.marketIndex, + jupiterClient: jupiterClient, + slippageBps: params.slippage ?? 100, + swapMode: "ExactIn", + }); + + txSig = signature; + } + + // @ts-expect-error - false undefined type conflict + if (params.toAmount) { + const jupiterClient = new JupiterClient({ connection: agent.connection }); + // @ts-expect-error - false undefined type conflict + const toAmount = numberToSafeBN(params.toAmount, toToken.precision); + const signature = await driftClient.swap({ + amount: toAmount, + inMarketIndex: toToken.marketIndex, + outMarketIndex: fromToken.marketIndex, + jupiterClient: jupiterClient, + slippageBps: params.slippage ?? 100, + swapMode: "ExactOut", + }); + + txSig = signature; + } + + await cleanUp(); + + // @ts-expect-error - false use before assignment + if (txSig) { + return txSig; + } + + throw new Error("Either fromAmount or toAmount must be provided"); + } catch (e) { + // @ts-expect-error error message is a string + throw new Error(`Failed to swap token: ${e.message}`); + } +} diff --git a/src/tools/drift/drift_vault.ts b/src/tools/drift/drift_vault.ts index feb23f1..c8d1214 100644 --- a/src/tools/drift/drift_vault.ts +++ b/src/tools/drift/drift_vault.ts @@ -37,14 +37,18 @@ export function getMarketIndexAndType(name: `${string}-${string}`) { if (type === "PERP") { const token = MainnetPerpMarkets.find((v) => v.baseAssetSymbol === symbol); if (!token) { - throw new Error("Drift doesn't have that market"); + throw new Error( + `Drift doesn't have that market. Here's a list of available perp markets: ${MainnetPerpMarkets.map((v) => v.baseAssetSymbol).join(", ")}`, + ); } return { marketIndex: token.marketIndex, marketType: MarketType.PERP }; } const token = MainnetSpotMarkets.find((v) => v.symbol === symbol); if (!token) { - throw new Error("Drift doesn't have that market"); + throw new Error( + `Drift doesn't have that market. Here's a list of available spot markets: ${MainnetSpotMarkets.map((v) => v.symbol).join(", ")}`, + ); } return { marketIndex: token.marketIndex, marketType: MarketType.SPOT }; } @@ -134,22 +138,22 @@ export async function createVault( const { vaultClient, driftClient, cleanUp } = await initClients(agent); const marketIndexAndType = getMarketIndexAndType(params.marketName); - if (!marketIndexAndType) { - throw new Error("Invalid market name"); - } - const spotMarket = driftClient.getSpotMarketAccount( marketIndexAndType.marketIndex, ); if (!spotMarket) { - throw new Error("Market not found"); + throw new Error( + `Market not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map((v) => `${v.symbol}-SPOT`).join(", ")}`, + ); } const spotPrecision = TEN.pow(new BN(spotMarket.decimals)); if (marketIndexAndType.marketType === MarketType.PERP) { - throw new Error("Only SPOT market names are supported"); + throw new Error( + `Only SPOT market names are supported. Such as ${MainnetSpotMarkets.map((v) => `${v.symbol}-SPOT`).join(", ")}`, + ); } const tx = await vaultClient.initializeVault({ @@ -239,7 +243,9 @@ export async function updateVault( ); if (!spotMarket) { - throw new Error("Market not found"); + throw new Error( + "Market not found. This vault's market is no longer supported", + ); } const spotPrecision = TEN.pow(new BN(spotMarket.decimals)); @@ -370,7 +376,9 @@ export async function depositIntoVault( ); if (!spotMarket) { - throw new Error("Market not found"); + throw new Error( + "Market not found. This vaults market is no longer supported", + ); } const spotPrecision = TEN.pow(new BN(spotMarket.decimals)); @@ -544,7 +552,7 @@ export async function tradeDriftVault( if (!isOwned) { throw new Error( - "This vault is owned by someone else, so you can't trade with it", + "This vault is owned/delegated to someone else, you can't trade with it", ); }