feat(drift): tools with params

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<typeof metadata.parameters>,
+        { 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,
+  };
+};

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<typeof metadata.parameters>,
+        { 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,
+  };
+};

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<string, ToolConfig> = {
   createDriftAccount: createDriftAccount(),
   getDriftAccountInfo: getDriftAccountInfo(),
   depositToDriftAccount: depositToDriftAccount(),
+  depositToDriftAccountWithParams: depositToDriftAccountWithParams(),
+  createDriftAccountWithParams: createDriftAccountWithParams(),
   getDriftAPY: getDriftAPY(),
   ...allMagicEdenTools,
   ...allChartTools,

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',