AI Tool Calling

import { tool } from "ai";
import { evaluate } from "mathjs";
import { z } from "zod";

export const weatherTool = tool({
  description: "Get the current weather in a location",
  inputSchema: z.object({
    location: z.string().describe("The city and country"),
    unit: z.enum(["celsius", "fahrenheit"]).default("celsius"),
  }),
  execute: async ({ location, unit }) => {
    // Replace with actual weather API call
    return {
      location,
      temperature: Math.round(Math.random() * 30),
      unit,
      conditions: "partly cloudy",
    };
  },
});

export const searchTool = tool({
  description: "Search the web for information",
  inputSchema: z.object({
    query: z.string().describe("The search query"),
    maxResults: z.number().optional().default(5),
  }),
  execute: async ({ query, maxResults }) => {
    // Replace with actual search API (Tavily, Brave, etc.)
    return {
      query,
      results: [{ title: `Result for: ${query}`, url: "https://example.com" }],
    };
  },
});

export const calculatorTool = tool({
  description: "Perform mathematical calculations",
  inputSchema: z.object({
    expression: z.string().describe("Math expression to evaluate"),
  }),
  execute: async ({ expression }) => {
    // mathjs.evaluate parses a math grammar instead of running JS, so
    // model/user-supplied input can't reach the runtime like eval would.
    try {
      return { expression, result: evaluate(expression) };
    } catch {
      return { expression, error: "Invalid expression" };
    }
  },
});

export const tools = {
  weather: weatherTool,
  search: searchTool,
  calculator: calculatorTool,
};

export type ToolName = keyof typeof tools;

import { streamText, generateText, stepCountIs, type ModelMessage } from "ai";
import { openai } from "@ai-sdk/openai";
import { tools, type ToolName } from "./tools";

// Uses OpenAI by default. Swap this for your own provider/model selection.
const model = openai("gpt-4o");

export interface AgentConfig {
  systemPrompt: string;
  availableTools: ToolName[];
  maxSteps?: number;
}

function pickTools(names: ToolName[]) {
  return Object.fromEntries(names.map((name) => [name, tools[name]] as const)) as Pick<
    typeof tools,
    ToolName
  >;
}

export async function runAgent({
  config,
  messages,
}: {
  config: AgentConfig;
  messages: ModelMessage[];
}) {
  const { systemPrompt, availableTools, maxSteps = 5 } = config;

  const result = await generateText({
    model,
    system: systemPrompt,
    messages,
    tools: pickTools(availableTools),
    stopWhen: stepCountIs(maxSteps),
  });

  return {
    text: result.text,
    toolCalls: result.toolCalls,
    toolResults: result.toolResults,
  };
}

export async function streamAgent({
  config,
  messages,
}: {
  config: AgentConfig;
  messages: ModelMessage[];
}) {
  const { systemPrompt, availableTools, maxSteps = 5 } = config;

  return streamText({
    model,
    system: systemPrompt,
    messages,
    tools: pickTools(availableTools),
    stopWhen: stepCountIs(maxSteps),
  });
}

export const agentConfigs: Record<string, AgentConfig> = {
  assistant: {
    systemPrompt: `You are a helpful AI assistant with access to various tools.
Use the available tools when needed to provide accurate responses.`,
    availableTools: ["weather", "search", "calculator"],
    maxSteps: 5,
  },
  researcher: {
    systemPrompt: `You are a research assistant. Use the search tool to find information.`,
    availableTools: ["search"],
    maxSteps: 5,
  },
};

import { NextRequest } from "next/server";
import { streamAgent, agentConfigs } from "@/lib/ai/agent";

export const runtime = "edge";

export async function POST(req: NextRequest) {
  const { messages, agentType = "assistant" } = await req.json();

  const config = agentConfigs[agentType];
  if (!config) {
    return Response.json({ error: "Unknown agent type" }, { status: 400 });
  }

  const result = await streamAgent({ config, messages });
  return result.toUIMessageStreamResponse();
}