diff --git a/src/ai/tools/drift/create-account-with-params.ts b/src/ai/tools/drift/create-account-with-params.ts new file mode 100644 index 0000000..cb696a4 --- /dev/null +++ b/src/ai/tools/drift/create-account-with-params.ts @@ -0,0 +1,83 @@ +import { tool } from 'ai'; +import { z } from 'zod'; + +import { streamUpdate } from '@/lib/utils'; +import { + createDriftAccount as createDriftAccountAction, +} from '@/server/actions/drift'; + +import { ToolConfig, WrappedToolProps } from '..'; + +export const createDriftAccountWithParams = (): ToolConfig => { + const metadata = { + description: 'Call this tool when the user wants to create a drift account (amount and symbol are required)', + parameters: z.object({ + amount: z.number().default(0).describe('The amount of tokens to deposit'), + symbol: z + .string() + .default('') + .describe('The symbol of the token to deposit'), + }), + }; + + const buildTool = ({ + dataStream = undefined, + abortData, + extraData: { agentKit }, + }: WrappedToolProps) => + tool({ + ...metadata, + execute: async ( + { amount, symbol }: z.infer, + { toolCallId }, + ) => { + if(!amount || !symbol) { + throw new Error('Amount and symbol are required'); + } + + streamUpdate({ + stream: dataStream, + update: { + type: 'stream-result-data', + toolCallId, + content: { + step: 'processing', + }, + }, + }); + + const result = await createDriftAccountAction({ + amount: amount, + symbol: symbol, + }); + + streamUpdate({ + stream: dataStream, + update: { + type: 'stream-result-data', + toolCallId, + status: 'idle', + content: { + ...result.result, + step: result.success ? 'completed' : 'failed', + }, + }, + }); + + return { + success: result.success, + result: { + ...result.result, + step: result.success ? 'completed' : 'failed', + }, + }; + + }, + }); + + return { + metadata, + buildTool, + confirm: createDriftAccountAction, + }; +}; diff --git a/src/ai/tools/drift/deposit-to-account-with-params.ts b/src/ai/tools/drift/deposit-to-account-with-params.ts new file mode 100644 index 0000000..e2c8ffd --- /dev/null +++ b/src/ai/tools/drift/deposit-to-account-with-params.ts @@ -0,0 +1,81 @@ +import { tool } from 'ai'; +import { z } from 'zod'; + +import { streamUpdate } from '@/lib/utils'; +import { + depositToDriftUserAccount, +} from '@/server/actions/drift'; + +import { ToolConfig, WrappedToolProps } from '..'; + +export const depositToDriftAccountWithParams = (): ToolConfig => { + const metadata = { + description: + 'Call this tool when the user wants to deposit to a drift account (amount and symbol are required)', + parameters: z.object({ + amount: z.number().optional().describe('The amount of tokens to deposit'), + symbol: z.string().optional().describe('The symbol of the token to deposit'), + }), + }; + + const buildTool = ({ + dataStream = undefined, + abortData, + extraData: { agentKit }, + }: WrappedToolProps) => + tool({ + ...metadata, + execute: async ( + { amount, symbol }: z.infer, + { toolCallId }, + ) => { + if(!amount || !symbol) { + throw new Error('Amount and symbol are required'); + } + + streamUpdate({ + stream: dataStream, + update: { + type: 'stream-result-data', + toolCallId, + content: { + step: 'processing', + }, + }, + }); + + const result = await depositToDriftUserAccount({ + amount: amount, + symbol: symbol, + }); + + streamUpdate({ + stream: dataStream, + update: { + type: 'stream-result-data', + toolCallId, + status: 'idle', + content: { + ...result.result, + step: result.success ? 'completed' : 'failed', + }, + }, + }); + + return { + success: result.success, + result: { + ...result.result, + step: result.success ? 'completed' : 'failed', + }, + }; + + }, + }); + + return { + metadata, + buildTool, + confirm: depositToDriftUserAccount, + }; +}; diff --git a/src/ai/tools/index.ts b/src/ai/tools/index.ts index c3cb94b..e1160e5 100644 --- a/src/ai/tools/index.ts +++ b/src/ai/tools/index.ts @@ -28,6 +28,8 @@ import { swapTokens } from './swap'; import { allTelegramTools } from './telegram'; import { transferTokens } from './transfer'; import { allTwitterTools } from './twitter'; +import { depositToDriftAccountWithParams } from './drift/deposit-to-account-with-params'; +import { createDriftAccountWithParams } from './drift/create-account-with-params'; interface ToolMetadata { description: string; @@ -66,6 +68,8 @@ export const allTools: Record = { createDriftAccount: createDriftAccount(), getDriftAccountInfo: getDriftAccountInfo(), depositToDriftAccount: depositToDriftAccount(), + depositToDriftAccountWithParams: depositToDriftAccountWithParams(), + createDriftAccountWithParams: createDriftAccountWithParams(), getDriftAPY: getDriftAPY(), ...allMagicEdenTools, ...allChartTools, diff --git a/src/app/(user)/chat/[id]/chat-interface.tsx b/src/app/(user)/chat/[id]/chat-interface.tsx index 46b9bea..0fdaa18 100644 --- a/src/app/(user)/chat/[id]/chat-interface.tsx +++ b/src/app/(user)/chat/[id]/chat-interface.tsx @@ -101,6 +101,14 @@ const TOOL_COMPONENTS: Record< component: DriftDeposit, displayName: '💰 Deposit to Drift', }, + depositToDriftAccountWithParams: { + component: DriftDeposit, + displayName: '💰 Deposit to Drift', + }, + createDriftAccountWithParams: { + component: DriftCard, + displayName: '👤 Create Drift Account', + }, tradeDriftPerpAccount: { component: DriftPrepTrade, displayName: '🌊 Prep trade',