Tool calling is how Athena moves from a plain chatbot into an agent. Instead of returning unstructured text, the model calls tools that return typed, structured data — project overviews, requirements, database schemas, and more. That data is then rendered as rich UI blocks in the chat.

Every tool in Athena follows the same four-step pattern:

Step 1: Zod Schema — define the shape of the tool's input
Step 2: TypeScript Types — type the data on the frontend
Step 3: Tool Definition — register the tool with the AI SDK
Step 4: API Route — expose the tool to the model

Once wired up, you render tool responses in the UI by switching on each message part's type.

Athena tools use input schemas only — the model fills in the schema and the result arrives on part.input. The execute function returns null because the model itself generates the structured output. You only need an outputSchema when the tool receives data back from the client (like the quiz answers tool).

How It Works

When a user sends a message, the flow looks like this:

The frontend sends the conversation to /api/chat
streamText passes the messages and registered tools to the model
The model decides which tool to call and fills in the schema
The response streams back as UI message parts — text parts and tool parts
The frontend maps over msg.parts and renders each tool's UI

The tool key you register in the API route becomes the part type in the UI. For example, registering projectOverview produces parts with type tool-projectOverview.

Step 1 Zod Schema

Start in lib/zodSchema.ts. Define a Zod object that describes exactly what structured data the model should return. Use .describe() on every field — these descriptions guide the model when filling in the schema.

Here is the schema for the project overview tool:

/lib/zodSchema.ts

export const ProjectOverviewSchema = z.object({
description: z.string().describe("Description of the project"),
goals_data: z.object({
  problem_statement: z.string().describe("Understood problem statement"),
  product_vision: z
    .string()
    .describe("The goal the vision behind the project"),
  goals: z.array(z.string()),
  success_metric: z.array(z.string()),
}),
});

For tools that need client-side input (like the quiz), define both an input and an output schema:

/lib/zodSchema.ts

export const AIQuestionsSchema = z.object({
questions: z.array(
  z.object({
    question: z.string().describe("Question about project to be answered"),
    options: z.array(
      z.string().describe("Answer helpers related to the question asked"),
    ),
    id: z.string().describe("Unique identifier for each question"),
  }),
),
});

export const AIQuestionsOutputSchema = z.object({
answers: z.array(
  z.object({
    id: z.string().describe("Unique id of the question answered"),
    answer: z.string().describe("Answer choosen by student"),
  }),
),
});

Step 2 TypeScript Types

Mirror your Zod schemas as TypeScript interfaces in lib/types.ts. These types are used when casting part.input in the UI so you get autocomplete and type safety.

/lib/types.ts

export interface ProjectOverview {
description: string;
goals_data: {
  problem_statement: string;
  product_vision: string;
  goals: string[];
  success_metric: string[];
};
}

export interface QuestionsPart {
input: {
  questions: Question[];
};
state:
  | "output-available"
  | "input-available"
  | "input-streaming"
  | "output-error";
}

Keep the interface fields in sync with the Zod schema. If you add a field to the schema, add it to the interface too.

Step 3 Tool Definition

In lib/tools.ts, import your schema and define the tool using the tool() helper from the AI SDK:

/lib/tools.ts

import { tool } from "ai";
import { ProjectOverviewSchema } from "./zodSchema";

export const projectOverviewGeneration = tool({
description:
  "This generates a clear description about what the project is about and it's goals.",
inputSchema: ProjectOverviewSchema,
execute: async () => {
  return null;
},
});

The description tells the model when to use the tool. Be specific — the model reads this to decide whether to call it. The inputSchema is the Zod schema from step 1.

For interactive tools that receive user input back, also pass an outputSchema:

/lib/tools.ts

export const askProjectQuestions = tool({
description:
  "This generate questions based on the project to improve project understanding",
outputSchema: AIQuestionsOutputSchema,
inputSchema: AIQuestionsSchema,
execute: async () => {
  return null;
},
});

Step 4 API Route

Register your tools in app/api/chat/route.ts inside the streamText call. The object key becomes the tool name the model sees and the prefix for UI part types.

/app/api/chat/route.ts

import {
askProjectQuestions,
projectOverviewGeneration,
functionalRequirementsGeneration,
userStoriesGeneration,
} from "@/lib/tools";

const result = streamText({
model: google("gemini-3.1-flash-lite-preview"),
messages: await convertToModelMessages(messages),
system: athena,
tools: {
  questionsTools: askProjectQuestions,
  projectOverview: projectOverviewGeneration,
  functionalRequirementsGeneration: functionalRequirementsGeneration,
  userStories: userStoriesGeneration,
},
stopWhen: stepCountIs(15),
});

stopWhen: stepCountIs(15) controls how many tool calls the model can make in a single request. Athena calls many tools in one run (overview, requirements, database design, etc.), so this needs to be high enough to allow the full planning flow.

Also reference your tools in the system prompt (lib/athena.ts) so the model knows which tool to use for each deliverable:

/lib/athena.ts

- User stories → use userStories
- Database design → use databaseDesign
- API routes and configuration → use apiDesign
- Tech stack required → use techStack
- Development plan with durations → use projectPlan
- Project risks with proposed solutions → use projectRisks

Rendering Tool UI

Tool results arrive inside messages from useChat(). Each message has a parts array. Loop over it and switch on part.type to render the right component.

The naming convention is tool-{toolKey} — if you registered projectOverview in the route, the part type is tool-projectOverview.

/app/athena/page.tsx

{messages.map((msg) => (
<div key={msg.id}>
  {msg.parts.map((part, i) => {
    switch (part.type) {
      case "text":
        return <span key={i}>{part.text}</span>;

      case "tool-projectOverview":
        const overview = part.input as ProjectOverview;
        if (!part.input || part.state !== "output-available") return null;

        return (
          <div key={i}>
            <h2>Project Overview</h2>
            <p>{overview.description}</p>
            <p>{overview.goals_data.problem_statement}</p>
          </div>
        );
    }
  })}
</div>
))}

Tool States

Every tool part has a state field. Always check it before rendering:

input-streaming — the model is still generating the tool input
input-available — input is ready but not yet confirmed
output-available — the tool input is complete, safe to render
output-error — something went wrong

In practice, most read-only display tools only render when part.state === "output-available" and part.input exists.

Interactive Tools

Some tools need user interaction before the conversation continues — like the quiz tool. The pattern is:

Render a custom component (e.g. QuestionsBlock) when the tool part is available
Collect user input in local state
Call sendMessage() with the collected answers when done

/app/athena/page.tsx

case "tool-questionsTools":
const toolpart = part as QuestionsPart;
if (!toolpart.input || toolpart.state !== "output-available") return null;

const handleAnswer = (id: string, option: string) => {
  const updatedAnswers = { ...answers, [id]: option };

  if (current === toolpart.input.questions.length - 1) {
    const summary = toolpart.input.questions
      .map((q) => `${q.question}: ${updatedAnswers[q.id]}`)
      .join("\n");
    sendMessage({ text: summary });
    setAnswers({});
    setCurrent(0);
  } else {
    setAnswers(updatedAnswers);
    setCurrent((prev) => prev + 1);
  }
};

return (
  <QuestionsBlock
    onClick={({ id, option }) => handleAnswer(id, option)}
    current={current}
    questions={toolpart.input.questions}
  />
);

When building the summary on the last question, use the updated answers object — not the stale React state. The current selection lives in updatedAnswers, but answers has not re-rendered yet.

Adding a New Tool

To add a new tool, work through the four steps in order:

Add a Zod schema in lib/zodSchema.ts
Add a matching interface in lib/types.ts
Export a tool() definition in lib/tools.ts
Register it in the tools object in app/api/chat/route.ts
Add a case "tool-yourToolKey": block in app/athena/page.tsx
Tell the model to use it in lib/athena.ts

That is the full loop — from schema to structured UI.