From 79fe5b0cb4eff7ae7039662164c8eccceae13b41 Mon Sep 17 00:00:00 2001 From: michaelessiet Date: Thu, 16 Jan 2025 21:33:07 +0100 Subject: [PATCH] fix: bugs noticed during testing --- .../requestUnstakeFromDriftInsuranceFund.ts | 3 +- src/actions/drift/swapSpotToken.ts | 10 +++-- .../drift/unstakeFromDriftInsuranceFund.ts | 2 +- src/tools/drift/drift.ts | 42 ++++++++++++++++++- src/vercel-ai/index.ts | 10 ++++- 5 files changed, 57 insertions(+), 10 deletions(-) diff --git a/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts b/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts index fc1b645..ca3ef0d 100644 --- a/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts +++ b/src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts @@ -8,8 +8,7 @@ const requestUnstakeFromDriftInsuranceFundAction: Action = { "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", + "request to unstake an amount from the drift insurance fund", ], examples: [ [ diff --git a/src/actions/drift/swapSpotToken.ts b/src/actions/drift/swapSpotToken.ts index a076cb8..1ca07e2 100644 --- a/src/actions/drift/swapSpotToken.ts +++ b/src/actions/drift/swapSpotToken.ts @@ -4,11 +4,13 @@ import { swapSpotToken } from "../../tools"; const driftSpotTokenSwapAction: Action = { name: "DRIFT_SPOT_TOKEN_SWAP_ACTION", - description: "Swap a spot token for another spot token on Drift", + description: "Swap a token for another 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", + "swap usdc to 5 sol on drift (in this case 5 sol is the toAmount)", + "swap 5 usdt to DRIFT on drift (in this case 5 usdt is the fromAmount)", ], examples: [ [ @@ -33,18 +35,18 @@ const driftSpotTokenSwapAction: Action = { fromAmount: z .number() .positive() - .describe("Amount to swap from in normal units e.g 50 === 50 SOL") + .describe("Amount to swap from e.g 50 === 50 SOL") .optional(), toAmount: z .number() .positive() - .describe("Amount to swap to in normal units e.g 5000 === 5000 USDC") + .describe("Amount to swap to e.g 5000 === 5000 USDC") .optional(), slippage: z .number() .positive() .describe("Slippage tolerance in percentage e.g 0.5 === 0.5%") - .optional(), + .default(0.5), }), handler: async (agent, input) => { try { diff --git a/src/actions/drift/unstakeFromDriftInsuranceFund.ts b/src/actions/drift/unstakeFromDriftInsuranceFund.ts index 4d9516e..2b11056 100644 --- a/src/actions/drift/unstakeFromDriftInsuranceFund.ts +++ b/src/actions/drift/unstakeFromDriftInsuranceFund.ts @@ -5,7 +5,7 @@ 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", + "Unstake requested unstake token from the Drift Insurance fund once the cool period has elapsed", similes: [ "unstake from the drift insurance fund", "withdraw from the drift insurance fund", diff --git a/src/tools/drift/drift.ts b/src/tools/drift/drift.ts index 57999f5..eb3f27c 100644 --- a/src/tools/drift/drift.ts +++ b/src/tools/drift/drift.ts @@ -4,6 +4,7 @@ import { DRIFT_PROGRAM_ID, DriftClient, FastSingleTxSender, + getInsuranceFundStakeAccountPublicKey, getLimitOrderParams, getMarketOrderParams, getUserAccountPublicKeySync, @@ -514,6 +515,25 @@ export async function stakeToDriftInsuranceFund( ); } + const deriveInsuranceFundStakeAccount = + getInsuranceFundStakeAccountPublicKey( + driftClient.program.programId, + agent.wallet.publicKey, + token.marketIndex, + ); + let shouldCreateAccount = false; + + try { + await driftClient.connection.getAccountInfo( + deriveInsuranceFundStakeAccount, + ); + } catch (e) { + // @ts-expect-error - error message is a string + if (e.message.includes("Account not found")) { + shouldCreateAccount = true; + } + } + const signature = await driftClient.addInsuranceFundStake({ amount: numberToSafeBN(amount, token.precision), marketIndex: token.marketIndex, @@ -521,6 +541,7 @@ export async function stakeToDriftInsuranceFund( token.mint, agent.wallet.publicKey, ), + initializeStakeAccount: shouldCreateAccount, txParams: { computeUnitsPrice: 0.000002 * 1000000 * 1000000, }, @@ -620,6 +641,7 @@ export async function unstakeFromDriftInsuranceFund( * @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 + * @param params.slippage slippage tolerance in percentage */ export async function swapSpotToken( agent: SolanaAgentKit, @@ -668,12 +690,20 @@ export async function swapSpotToken( const jupiterClient = new JupiterClient({ connection: agent.connection }); // @ts-expect-error - false undefined type conflict const fromAmount = numberToSafeBN(params.fromAmount, fromToken.precision); + const res = await ( + await fetch( + `https://quote-api.jup.ag/v6/quote?inputMint=${fromToken.mint}&outputMint=${toToken.mint}&amount=${fromAmount.toNumber()}&slippageBps=${(params.slippage ?? 0.5) * 100}&swapMode=ExactIn`, + ) + ).json(); const signature = await driftClient.swap({ amount: fromAmount, inMarketIndex: fromToken.marketIndex, outMarketIndex: toToken.marketIndex, jupiterClient: jupiterClient, - slippageBps: params.slippage ?? 100, + v6: { + quote: res, + }, + slippageBps: (params.slippage ?? 0.5) * 100, swapMode: "ExactIn", }); @@ -685,12 +715,20 @@ export async function swapSpotToken( const jupiterClient = new JupiterClient({ connection: agent.connection }); // @ts-expect-error - false undefined type conflict const toAmount = numberToSafeBN(params.toAmount, toToken.precision); + const res = await ( + await fetch( + `https://quote-api.jup.ag/v6/quote?inputMint=${fromToken.mint}&outputMint=${toToken.mint}&amount=${toAmount.toNumber()}&slippageBps=${(params.slippage ?? 0.5) * 100}&swapMode=ExactOut`, + ) + ).json(); const signature = await driftClient.swap({ amount: toAmount, inMarketIndex: toToken.marketIndex, outMarketIndex: fromToken.marketIndex, jupiterClient: jupiterClient, - slippageBps: params.slippage ?? 100, + v6: { + quote: res, + }, + slippageBps: (params.slippage ?? 0.5) * 100, swapMode: "ExactOut", }); diff --git a/src/vercel-ai/index.ts b/src/vercel-ai/index.ts index ba35643..8d1219a 100644 --- a/src/vercel-ai/index.ts +++ b/src/vercel-ai/index.ts @@ -14,7 +14,15 @@ export function createSolanaTools( tools[key] = tool({ // @ts-expect-error Value matches type however TS still shows error id: action.name, - description: action.description, + description: ` + ${action.description} + + Similes: ${action.similes.map( + (simile) => ` + ${simile} + `, + )} + `.slice(0, 1023), parameters: action.schema, execute: async (params) => await executeAction(action, solanaAgentKit, params),