Custom AI Tools

import { tool } from "ai";
import { z } from "zod";
 
const myTool = tool({
  description: "A description of what this tool does",
  parameters: z.object({
    param1: z.string().describe("Description of param1"),
    param2: z.number().optional().describe("Optional param2"),
  }),
  execute: async ({ param1, param2 }) => {
    // Tool logic here
    return { result: "..." };
  },
});

// lib/ai/tools/weather.ts
import { tool } from "ai";
import { z } from "zod";
 
export const weatherTool = tool({
  description: "Get current weather for a location",
  parameters: z.object({
    location: z.string().describe("City name, e.g., 'San Francisco'"),
    units: z.enum(["celsius", "fahrenheit"]).default("celsius"),
  }),
  execute: async ({ location, units }) => {
    // Call weather API
    const response = await fetch(
      `https://api.weatherapi.com/v1/current.json?key=${process.env.WEATHER_API_KEY}&q=${location}`
    );
    const data = await response.json();
 
    const temp = units === "celsius" ? data.current.temp_c : data.current.temp_f;
 
    return {
      location: data.location.name,
      country: data.location.country,
      temperature: temp,
      units,
      condition: data.current.condition.text,
      humidity: data.current.humidity,
    };
  },
});

// lib/ai/tools/search.ts
import { tool } from "ai";
import { z } from "zod";
 
export const searchTool = tool({
  description: "Search the web for information",
  parameters: z.object({
    query: z.string().describe("The search query"),
    maxResults: z.number().default(5).describe("Maximum results to return"),
  }),
  execute: async ({ query, maxResults }) => {
    // Using Serper API as example
    const response = await fetch("https://google.serper.dev/search", {
      method: "POST",
      headers: {
        "X-API-KEY": process.env.SERPER_API_KEY!,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({ q: query, num: maxResults }),
    });
 
    const data = await response.json();
 
    return {
      results: data.organic.map((item: any) => ({
        title: item.title,
        link: item.link,
        snippet: item.snippet,
      })),
    };
  },
});

// lib/ai/tools/database.ts
import { tool } from "ai";
import { z } from "zod";
import { db } from "@/lib/db/client";
import { users, subscriptions } from "@/lib/db/schema";
import { eq, sql } from "drizzle-orm";
 
export const userStatsTool = tool({
  description: "Get statistics about users in the system",
  parameters: z.object({
    metric: z.enum(["total", "by_plan", "recent_signups"]),
  }),
  execute: async ({ metric }) => {
    switch (metric) {
      case "total":
        const [{ count }] = await db.select({ count: sql<number>`count(*)` }).from(users);
        return { totalUsers: count };
 
      case "by_plan":
        const byPlan = await db
          .select({
            plan: users.plan,
            count: sql<number>`count(*)`,
          })
          .from(users)
          .groupBy(users.plan);
        return { usersByPlan: byPlan };
 
      case "recent_signups":
        const recentUsers = await db
          .select({ name: users.name, createdAt: users.createdAt })
          .from(users)
          .orderBy(sql`${users.createdAt} DESC`)
          .limit(5);
        return { recentSignups: recentUsers };
    }
  },
});

// lib/ai/tools/calculator.ts
import { tool } from "ai";
import { z } from "zod";
 
export const calculatorTool = tool({
  description: "Perform mathematical calculations",
  parameters: z.object({
    expression: z.string().describe("Mathematical expression to evaluate, e.g., '2 + 2 * 3'"),
  }),
  execute: async ({ expression }) => {
    // Use a safe math parser in production
    // For demo, using Function (not recommended for production)
    try {
      // Sanitize - only allow numbers and operators
      const sanitized = expression.replace(/[^0-9+\-*/().%\s]/g, "");
      const result = Function(`return ${sanitized}`)();
      return { expression, result };
    } catch (error) {
      return { expression, error: "Invalid expression" };
    }
  },
});

// lib/ai/tools/file-reader.ts
import { tool } from "ai";
import { z } from "zod";
import { getDownloadUrl } from "@/lib/storage/r2";
 
export const readFileTool = tool({
  description: "Read content from a file stored in the system",
  parameters: z.object({
    fileKey: z.string().describe("The file key/path in storage"),
  }),
  execute: async ({ fileKey }) => {
    const url = await getDownloadUrl(fileKey);
    const response = await fetch(url);
    const content = await response.text();
 
    return {
      fileKey,
      contentLength: content.length,
      preview: content.slice(0, 1000),
    };
  },
});

// lib/ai/tools/index.ts
import { weatherTool } from "./weather";
import { searchTool } from "./search";
import { calculatorTool } from "./calculator";
import { userStatsTool } from "./database";
 
export const tools = {
  weather: weatherTool,
  search: searchTool,
  calculator: calculatorTool,
  userStats: userStatsTool,
};
 
export type ToolName = keyof typeof tools;

// app/api/chat/route.ts
import { streamText } from "ai";
import { models } from "@/lib/ai/models";
import { tools } from "@/lib/ai/tools";
 
export async function POST(req: Request) {
  const { messages, model = "gpt-4o" } = await req.json();
 
  const result = await streamText({
    model: models[model],
    messages,
    tools,
    maxToolRoundtrips: 5, // Allow multiple tool calls
    system: `You are a helpful assistant with access to various tools.
Use the appropriate tool when the user asks for:
- Weather information
- Web searches
- Calculations
- User statistics`,
  });
 
  return result.toDataStreamResponse();
}

// lib/ai/tools/user-context.ts
import { tool } from "ai";
import { z } from "zod";
import { db } from "@/lib/db/client";
import { users, subscriptions } from "@/lib/db/schema";
import { eq } from "drizzle-orm";
 
export function createUserContextTool(userId: string) {
  return tool({
    description: "Get information about the current user",
    parameters: z.object({
      info: z.enum(["profile", "subscription", "usage"]),
    }),
    execute: async ({ info }) => {
      switch (info) {
        case "profile":
          const [user] = await db.select().from(users).where(eq(users.id, userId));
          return { profile: user };
 
        case "subscription":
          const [sub] = await db
            .select()
            .from(subscriptions)
            .where(eq(subscriptions.userId, userId));
          return { subscription: sub };
 
        case "usage":
          // Get usage stats
          return {
            usage: {
              /* ... */
            },
          };
      }
    },
  });
}
 
// Usage in API route
const result = await streamText({
  model: models[model],
  messages,
  tools: {
    ...tools,
    userContext: createUserContextTool(session.user.id),
  },
});

// lib/ai/tools/async-task.ts
import { tool } from "ai";
import { z } from "zod";
import { longRunningTask } from "@/trigger/tasks";
 
export const asyncTaskTool = tool({
  description: "Start a long-running task",
  parameters: z.object({
    taskType: z.string(),
    data: z.record(z.any()),
  }),
  execute: async ({ taskType, data }) => {
    // Trigger background job
    const handle = await longRunningTask.trigger({ taskType, data });
 
    return {
      taskId: handle.id,
      message: "Task started. Check back later for results.",
    };
  },
});

// For sensitive operations, return a confirmation step
export const deleteDataTool = tool({
  description: "Delete user data (requires confirmation)",
  parameters: z.object({
    dataType: z.enum(["chats", "files", "account"]),
    confirmed: z.boolean().default(false),
  }),
  execute: async ({ dataType, confirmed }) => {
    if (!confirmed) {
      return {
        requiresConfirmation: true,
        message: `Are you sure you want to delete your ${dataType}? This cannot be undone.`,
        confirmationPrompt: `To confirm, say "delete my ${dataType}"`,
      };
    }
 
    // Perform deletion
    return { success: true, message: `${dataType} deleted` };
  },
});

// components/chat/tool-result.tsx
import { Card } from "@/components/ui/card";
import { Cloud, Search, Calculator } from "lucide-react";
 
interface ToolResultProps {
  toolName: string;
  result: any;
}
 
export function ToolResult({ toolName, result }: ToolResultProps) {
  switch (toolName) {
    case "weather":
      return (
        <Card className="p-4">
          <div className="flex items-center gap-2">
            <Cloud className="h-5 w-5" />
            <span className="font-medium">{result.location}</span>
          </div>
          <p className="mt-2 text-2xl">
            {result.temperature}° {result.units === "celsius" ? "C" : "F"}
          </p>
          <p className="text-neutral-500 dark:text-neutral-400">{result.condition}</p>
        </Card>
      );
 
    case "search":
      return (
        <Card className="p-4">
          <div className="mb-3 flex items-center gap-2">
            <Search className="h-5 w-5" />
            <span className="font-medium">Search Results</span>
          </div>
          <ul className="space-y-2">
            {result.results.map((r: any, i: number) => (
              <li key={i}>
                <a href={r.link} className="text-emerald-600 hover:underline">
                  {r.title}
                </a>
                <p className="text-sm text-neutral-500 dark:text-neutral-400">{r.snippet}</p>
              </li>
            ))}
          </ul>
        </Card>
      );
 
    default:
      return (
        <Card className="p-4">
          <pre className="text-sm">{JSON.stringify(result, null, 2)}</pre>
        </Card>
      );
  }
}