feat: add more drift actions

This commit is contained in:
michaelessiet
2025-01-17 17:12:34 +01:00
parent 79fe5b0cb4
commit 0c840d9bcb
12 changed files with 625 additions and 1 deletions

View File

@@ -0,0 +1,65 @@
import { z } from "zod";
import type { Action } from "../../types";
import { getEntryQuoteOfPerpTrade } from "../../tools";
const entryQuoteOfPerpTradeAction: Action = {
name: "DRIFT_GET_ENTRY_QUOTE_OF_PERP_TRADE_ACTION",
description: "Get the entry quote of a perpetual trade on Drift",
similes: [
"get the entry quote of a perpetual trade on drift",
"get the entry quote of a perp trade on drift",
"get the entry quote of the BTC-PERP trade on drift",
"get the entry quote of the SOL-PERP trade on drift",
"get the entry quote of a 1000 USDC long on the SOL-PERP market",
"get the entry quote of a 1000 USDC short on the SOL-PERP market",
"quote for a $1000 long on the BTC-PERP market",
],
examples: [
[
{
input: {
marketSymbol: "BTC-PERP",
type: "long",
amount: 1000,
},
output: {
status: "success",
data: {
entryPrice: 100000,
priceImpact: 0.0001,
bestPrice: 100001,
worstPrice: 99999,
baseFilled: 1000,
quoteFilled: 1000,
},
},
explanation:
"Get the entry quote of a $1000 long on the BTC-PERP market",
},
],
],
schema: z.object({
marketSymbol: z.string().describe("Symbol of the perpetual market"),
type: z.enum(["long", "short"]).describe("Type of trade"),
amount: z.number().positive().describe("Amount to trade"),
}),
handler: async (agent, input) => {
try {
const data = await getEntryQuoteOfPerpTrade(
input.marketSymbol,
input.amount,
input.type,
);
return data;
} catch (e) {
return {
status: "error",
// @ts-expect-error error is not a string
message: e.message,
};
}
},
};
export default entryQuoteOfPerpTradeAction;

View File

@@ -0,0 +1,52 @@
import { z } from "zod";
import type { Action } from "../../types";
import { getLendingAndBorrowAPY } from "../../tools";
const lendAndBorrowAPYAction: Action = {
name: "DRIFT_GET_LEND_AND_BORROW_APY_ACTION",
description: "Get the lending and borrowing APY (in %) of a token on Drift",
similes: [
"get the lending and borrowing APY of a token on drift",
"get the lending and borrowing APY of a token on drift",
"get the lending and borrowing APY of the USDC token on drift",
"get the lending and borrowing APY of the SOL token on drift",
],
examples: [
[
{
input: {
symbol: "USDC",
},
output: {
status: "success",
data: {
lendingAPY: 10,
borrowingAPY: 12.1,
},
},
explanation: "Get the lending and borrowing APY of the USDC token",
},
],
],
schema: z.object({
symbol: z.string().describe("Symbol of the token"),
}),
handler: async (agent, input) => {
try {
const data = await getLendingAndBorrowAPY(agent, input.symbol);
return {
status: "success",
data,
};
} catch (e) {
return {
status: "error",
// @ts-expect-error error is not a string
message: e.message,
};
}
},
};
export default lendAndBorrowAPYAction;

View File

@@ -0,0 +1,61 @@
import { z } from "zod";
import type { Action } from "../../types";
import { calculatePerpMarketFundingRate } from "../../tools";
const perpMarktetFundingRateAction: Action = {
name: "DRIFT_PERP_MARKET_FUNDING_RATE_ACTION",
description: "Get the funding rate of a perpetual market on Drift",
similes: [
"get the yearly funding rate of a perpetual market on drift",
"get the funding rate of a perp market on drift",
"get the hourly funding rate of a perpetual market on drift",
"get the funding rate of the BTC-PERP market on drift",
"get the funding rate of the SOL-PERP market on drift",
],
examples: [
[
{
input: {
marketSymbol: "BTC-PERP",
},
output: {
status: "success",
data: {
longRate: 0.0001,
shortRate: 0.0002,
},
},
explanation: "Get the funding rate of the BTC-PERP market",
},
],
],
schema: z.object({
marketSymbol: z
.string()
.toUpperCase()
.describe("Symbol of the perpetual market"),
period: z.enum(["year", "hour"]).default("hour").optional(),
}),
handler: async (agent, input) => {
try {
const data = await calculatePerpMarketFundingRate(
agent,
input.marketSymbol,
input.period,
);
return {
status: "success",
data,
};
} catch (e) {
return {
status: "error",
// @ts-expect-error error is not a string
message: e.message,
};
}
},
};
export default perpMarktetFundingRateAction;

View File

@@ -64,6 +64,9 @@ import stakeToDriftInsuranceFundAction from "./drift/stakeToDriftInsuranceFund";
import requestUnstakeFromDriftInsuranceFundAction from "./drift/requestUnstakeFromDriftInsuranceFund";
import unstakeFromDriftInsuranceFundAction from "./drift/unstakeFromDriftInsuranceFund";
import driftSpotTokenSwapAction from "./drift/swapSpotToken";
import perpMarktetFundingRateAction from "./drift/perpMarketFundingRate";
import entryQuoteOfPerpTradeAction from "./drift/entryQuoteOfPerpTrade";
import lendAndBorrowAPYAction from "./drift/getLendAndBorrowAPY";
export const ACTIONS = {
WALLET_ADDRESS_ACTION: getWalletAddressAction,
@@ -134,6 +137,9 @@ export const ACTIONS = {
requestUnstakeFromDriftInsuranceFundAction,
UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION: unstakeFromDriftInsuranceFundAction,
DRIFT_SPOT_TOKEN_SWAP_ACTION: driftSpotTokenSwapAction,
DRIFT_PERP_MARKET_FUNDING_RATE_ACTION: perpMarktetFundingRateAction,
DRIFT_GET_ENTRY_QUOTE_OF_PERP_TRADE_ACTION: entryQuoteOfPerpTradeAction,
DRIFT_GET_LEND_AND_BORROW_APY_ACTION: lendAndBorrowAPYAction,
};
export type { Action, ActionExample, Handler } from "../types/action";

View File

@@ -104,6 +104,9 @@ import {
requestUnstakeFromDriftInsuranceFund,
unstakeFromDriftInsuranceFund,
swapSpotToken,
calculatePerpMarketFundingRate,
getEntryQuoteOfPerpTrade,
getLendingAndBorrowAPY,
} from "../tools";
import {
Config,
@@ -871,4 +874,20 @@ export class SolanaAgentKit {
slippage: params.slippage,
});
}
async getPerpMarketFundingRate(
symbol: `${string}-PERP`,
period: "year" | "hour" = "year",
) {
return calculatePerpMarketFundingRate(this, symbol, period);
}
async getEntryQuoteOfPerpTrade(
amount: number,
symbol: `${string}-PERP`,
action: "short" | "long",
) {
return getEntryQuoteOfPerpTrade(symbol, amount, action);
}
async getLendAndBorrowAPY(symbol: string) {
return getLendingAndBorrowAPY(this, symbol);
}
}

View File

@@ -0,0 +1,39 @@
import { Tool } from "langchain/tools";
import type { SolanaAgentKit } from "../../agent";
export class SolanaDriftEntryQuoteOfPerpTradeTool extends Tool {
name = "drift_entry_quote_of_perp_trade";
description = `Get an entry quote for a perpetual trade on Drift protocol.
Inputs (JSON string):
- amount: number, amount to trade (required)
- symbol: string, market symbol (required)
- action: "long" | "short", trade direction (required)`;
constructor(private solanaKit: SolanaAgentKit) {
super();
}
protected async _call(input: string): Promise<string> {
try {
const parsedInput = JSON.parse(input);
const quote = await this.solanaKit.getEntryQuoteOfPerpTrade(
parsedInput.amount,
parsedInput.symbol,
parsedInput.action,
);
return JSON.stringify({
status: "success",
message: `Entry quote retrieved for ${parsedInput.action} ${parsedInput.amount} ${parsedInput.symbol}`,
data: quote,
});
} catch (error: any) {
return JSON.stringify({
status: "error",
message: error.message,
code: error.code || "ENTRY_QUOTE_OF_PERP_TRADE_ERROR",
});
}
}
}

View File

@@ -13,3 +13,10 @@ export * from "./update_vault";
export * from "./vault_info";
export * from "./withdraw_from_account";
export * from "./withdraw_from_vault";
export * from "./perp_market_funding_rate";
export * from "./entry_quote_of_perp_trade";
export * from "./lend_and_borrow_apy";
export * from "./stake_to_insurance_fund";
export * from "./swap_spot_token";
export * from "./unstake_from_insurance_fund";
export * from "./request_unstake_from_insurance_fund";

View File

@@ -0,0 +1,32 @@
import { Tool } from "langchain/tools";
import type { SolanaAgentKit } from "../../agent";
export class SolanaDriftLendAndBorrowAPYTool extends Tool {
name = "drift_lend_and_borrow_apy";
description = `Get lending and borrowing APY for a token on Drift protocol.
Inputs (JSON string):
- symbol: string, token symbol (required)`;
constructor(private solanaKit: SolanaAgentKit) {
super();
}
protected async _call(input: string): Promise<string> {
try {
const apyInfo = await this.solanaKit.getLendAndBorrowAPY(input);
return JSON.stringify({
status: "success",
message: `APY information retrieved for ${input}`,
data: apyInfo,
});
} catch (error: any) {
return JSON.stringify({
status: "error",
message: error.message,
code: error.code || "LEND_AND_BORROW_APY_ERROR",
});
}
}
}

View File

@@ -0,0 +1,36 @@
import { Tool } from "langchain/tools";
import type { SolanaAgentKit } from "../../agent";
export class SolanaDriftPerpMarketFundingRateTool extends Tool {
name = "drift_perp_market_funding_rate";
description = `Get the funding rate for a perpetual market on Drift protocol.
Inputs (JSON string):
- symbol: string, market symbol (required)
- period: year or hour (default: hour)`;
constructor(private solanaKit: SolanaAgentKit) {
super();
}
protected async _call(input: string): Promise<string> {
try {
const parsedInput = JSON.parse(input);
const fundingRate = await this.solanaKit.getPerpMarketFundingRate(
parsedInput.symbol,
parsedInput.period,
);
return JSON.stringify({
status: "success",
message: `Funding rate retrieved for ${parsedInput.symbol}`,
data: fundingRate,
});
} catch (error: any) {
return JSON.stringify({
status: "error",
message: error.message,
});
}
}
}

View File

@@ -27,7 +27,7 @@ export * from "./squads";
export * from "./helius";
export * from "./drift";
import { SolanaAgentKit } from "../agent";
import type { SolanaAgentKit } from "../agent";
import {
SolanaBalanceTool,
SolanaBalanceOtherTool,
@@ -114,6 +114,13 @@ import {
SolanaUpdateDriftVaultTool,
SolanaWithdrawFromDriftAccountTool,
SolanaWithdrawFromDriftVaultTool,
SolanaDriftLendAndBorrowAPYTool,
SolanaDriftEntryQuoteOfPerpTradeTool,
SolanaDriftPerpMarketFundingRateTool,
SolanaDriftSpotTokenSwapTool,
SolanaRequestUnstakeFromDriftInsuranceFundTool,
SolanaStakeToDriftInsuranceFundTool,
SolanaUnstakeFromDriftInsuranceFundTool,
} from "./index";
export function createSolanaTools(solanaKit: SolanaAgentKit) {
@@ -208,5 +215,12 @@ export function createSolanaTools(solanaKit: SolanaAgentKit) {
new SolanaDriftVaultInfoTool(solanaKit),
new SolanaWithdrawFromDriftAccountTool(solanaKit),
new SolanaWithdrawFromDriftVaultTool(solanaKit),
new SolanaDriftSpotTokenSwapTool(solanaKit),
new SolanaStakeToDriftInsuranceFundTool(solanaKit),
new SolanaRequestUnstakeFromDriftInsuranceFundTool(solanaKit),
new SolanaUnstakeFromDriftInsuranceFundTool(solanaKit),
new SolanaDriftLendAndBorrowAPYTool(solanaKit),
new SolanaDriftEntryQuoteOfPerpTradeTool(solanaKit),
new SolanaDriftPerpMarketFundingRateTool(solanaKit),
];
}

View File

@@ -1,9 +1,16 @@
import {
BASE_PRECISION,
BigNum,
calculateDepositRate,
calculateEstimatedEntryPriceWithL2,
calculateInterestRate,
calculateLongShortFundingRateAndLiveTwaps,
convertToNumber,
DRIFT_PROGRAM_ID,
DriftClient,
FastSingleTxSender,
FUNDING_RATE_BUFFER_PRECISION,
FUNDING_RATE_PRECISION_EXP,
getInsuranceFundStakeAccountPublicKey,
getLimitOrderParams,
getMarketOrderParams,
@@ -12,6 +19,7 @@ import {
MainnetPerpMarkets,
MainnetSpotMarkets,
numberToSafeBN,
PERCENTAGE_PRECISION,
PositionDirection,
PostOnlyParams,
PRICE_PRECISION,
@@ -26,6 +34,7 @@ import { getAssociatedTokenAddressSync } from "@solana/spl-token";
import { PublicKey } from "@solana/web3.js";
import { Transaction } from "@solana/web3.js";
import { ComputeBudgetProgram } from "@solana/web3.js";
import type { RawL2Output } from "./types";
export async function initClients(
agent: SolanaAgentKit,
@@ -748,3 +757,254 @@ export async function swapSpotToken(
throw new Error(`Failed to swap token: ${e.message}`);
}
}
/**
* To get funding rate as a percentage, you need to multiply by the funding rate buffer precision
* @param rawFundingRate
*/
export function getFundingRateAsPercentage(rawFundingRate: anchor.BN) {
return BigNum.from(
rawFundingRate.mul(FUNDING_RATE_BUFFER_PRECISION),
FUNDING_RATE_PRECISION_EXP,
).toNum();
}
/**
* Calculate the funding rate for a perpetual market
* @param agent
* @param marketSymbol
*/
export async function calculatePerpMarketFundingRate(
agent: SolanaAgentKit,
marketSymbol: `${string}-PERP`,
period: "year" | "hour",
) {
try {
const { driftClient, cleanUp } = await initClients(agent);
const market = driftClient.getMarketIndexAndType(
`${marketSymbol.toUpperCase()}`,
);
if (!market) {
throw new Error(
`This market isn't available on the Drift Protocol. Here's a list of markets that are: ${MainnetPerpMarkets.map(
(v) => v.symbol,
).join(", ")}`,
);
}
const marketAccount = driftClient.getPerpMarketAccount(market.marketIndex);
if (!marketAccount) {
throw new Error("Market account not found");
}
const [
_marketTwapLive,
_oracleTwapLive,
longFundingRate,
shortFundingRate,
] = await calculateLongShortFundingRateAndLiveTwaps(
marketAccount,
driftClient.getOracleDataForPerpMarket(market.marketIndex),
undefined,
new anchor.BN(Date.now()),
);
await cleanUp();
let longFundingRateNum = getFundingRateAsPercentage(longFundingRate);
let shortFundingRateNum = getFundingRateAsPercentage(shortFundingRate);
if (period === "year") {
const paymentsPerYear = 24 * 365.25;
longFundingRateNum *= paymentsPerYear;
shortFundingRateNum *= paymentsPerYear;
}
const longsArePaying = longFundingRateNum > 0;
const shortsArePaying = !(shortFundingRateNum > 0);
const longsAreString = longsArePaying ? "pay" : "receive";
const shortsAreString = !shortsArePaying ? "receive" : "pay";
const absoluteLongFundingRateNum = Math.abs(longFundingRateNum);
const absoluteShortFundingRateNum = Math.abs(shortFundingRateNum);
const formattedLongRatePct = absoluteLongFundingRateNum.toFixed(
period === "hour" ? 5 : 2,
);
const formattedShortRatePct = absoluteShortFundingRateNum.toFixed(
period === "hour" ? 5 : 2,
);
const paymentUnit = period === "year" ? "% APR" : "%";
const friendlyString = `At this rate, longs would ${longsAreString} ${formattedLongRatePct} ${paymentUnit} and shorts would ${shortsAreString} ${formattedShortRatePct} ${paymentUnit} at the end of the hour.`;
return {
longRate: longsArePaying
? -absoluteLongFundingRateNum
: absoluteLongFundingRateNum,
shortRate: shortsArePaying
? -absoluteShortFundingRateNum
: absoluteShortFundingRateNum,
friendlyString,
};
} catch (e) {
throw new Error(
// @ts-expect-error e.message is a string
`Something went wrong while trying to get the market's funding rate. Here's some more context: ${e.message}`,
);
}
}
export async function getL2OrderBook(marketSymbol: `${string}-PERP`) {
try {
const serializedOrderbook: RawL2Output = await (
await fetch(
`https://dlob.drift.trade/l2?marketName=${marketSymbol.toUpperCase()}&includeOracle=true`,
)
).json();
return {
asks: serializedOrderbook.asks.map((ask) => ({
price: new anchor.BN(ask.price),
size: new anchor.BN(ask.size),
sources: Object.entries(ask.sources).reduce((previous, [key, val]) => {
return {
...(previous ?? {}),
[key]: new anchor.BN(val),
};
}, {}),
})),
bids: serializedOrderbook.bids.map((bid) => ({
price: new anchor.BN(bid.price),
size: new anchor.BN(bid.size),
sources: Object.entries(bid.sources).reduce((previous, [key, val]) => {
return {
...(previous ?? {}),
[key]: new anchor.BN(val),
};
}, {}),
})),
oracleData: {
price: serializedOrderbook.oracleData.price
? new anchor.BN(serializedOrderbook.oracleData.price)
: undefined,
slot: serializedOrderbook.oracleData.slot
? new anchor.BN(serializedOrderbook.oracleData.slot)
: undefined,
confidence: serializedOrderbook.oracleData.confidence
? new anchor.BN(serializedOrderbook.oracleData.confidence)
: undefined,
hasSufficientNumberOfDataPoints:
serializedOrderbook.oracleData.hasSufficientNumberOfDataPoints,
twap: serializedOrderbook.oracleData.twap
? new anchor.BN(serializedOrderbook.oracleData.twap)
: undefined,
twapConfidence: serializedOrderbook.oracleData.twapConfidence
? new anchor.BN(serializedOrderbook.oracleData.twapConfidence)
: undefined,
maxPrice: serializedOrderbook.oracleData.maxPrice
? new anchor.BN(serializedOrderbook.oracleData.maxPrice)
: undefined,
},
slot: serializedOrderbook.slot,
};
} catch (e) {
throw new Error();
}
}
/**
* Get the estimated entry quote of a perp trade
* @param agent
* @param marketSymbol
* @param amount
* @param type
*/
export async function getEntryQuoteOfPerpTrade(
marketSymbol: `${string}-PERP`,
amount: number,
type: "long" | "short",
) {
try {
const l2OrderBookData = await getL2OrderBook(marketSymbol);
const estimatedEntryPriceData = calculateEstimatedEntryPriceWithL2(
"quote",
numberToSafeBN(amount, BASE_PRECISION),
type === "long" ? PositionDirection.LONG : PositionDirection.SHORT,
BASE_PRECISION,
// @ts-expect-error - false type conflict
l2OrderBookData,
);
return {
entryPrice: convertToNumber(
estimatedEntryPriceData.entryPrice,
QUOTE_PRECISION,
),
priceImpact: convertToNumber(
estimatedEntryPriceData.priceImpact,
QUOTE_PRECISION,
),
bestPrice: convertToNumber(
estimatedEntryPriceData.bestPrice,
QUOTE_PRECISION,
),
worstPrice: convertToNumber(
estimatedEntryPriceData.worstPrice,
QUOTE_PRECISION,
),
};
} catch (e) {
// @ts-expect-error - error message is a string
throw new Error(`Failed to get entry quote: ${e.message}`);
}
}
/**
* Get the APY for lending and borrowing a specific token on drift protocol
* @param agent
* @param symbol
*/
export async function getLendingAndBorrowAPY(
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 marketAccount = driftClient.getSpotMarketAccount(token.marketIndex);
if (!marketAccount) {
throw new Error("Market account not found");
}
const lendAPY = calculateDepositRate(marketAccount);
const borrowAPY = calculateInterestRate(marketAccount);
await cleanUp();
return {
lendingAPY: convertToNumber(lendAPY, PERCENTAGE_PRECISION) * 100, // convert to percentage
borrowAPY: convertToNumber(borrowAPY, PERCENTAGE_PRECISION) * 100, // convert to percentage
};
} catch (e) {
// @ts-expect-error - error message is a string
throw new Error(`Failed to get APYs: ${e.message}`);
}
}

33
src/tools/drift/types.ts Normal file
View File

@@ -0,0 +1,33 @@
import type { L2OrderBook, MarketType, OraclePriceData } from "@drift-labs/sdk";
export type L2WithOracle = L2OrderBook & { oracleData: OraclePriceData };
export type RawL2Output = {
marketIndex: number;
marketType: MarketType;
marketName: string;
asks: {
price: string;
size: string;
sources: {
[key: string]: string;
};
}[];
bids: {
price: string;
size: string;
sources: {
[key: string]: string;
};
}[];
oracleData: {
price: string;
slot: string;
confidence: string;
hasSufficientNumberOfDataPoints: boolean;
twap?: string;
twapConfidence?: string;
maxPrice?: string;
};
slot?: number;
};