Adding Custom AI Tools
Relevant source files
The following files were used as context for generating this wiki page:
AI Tools extend the capabilities of the LV.AI assistant by allowing it to perform structured actions such as creating documents, updating artifacts, or requesting suggestions. Tools are defined on the server using the Vercel AI SDK and mapped to UI components for rendering within the chat interface.
Sources: lib/ai/tools/create-document.ts, lib/types.ts:24-31
Tool Architecture
Tools consist of a server-side definition and a corresponding client-side UI implementation. The server definition specifies the input schema using Zod and handles the execution logic, while the client side renders the tool's state and results.
The diagram above shows how a tool call originates from the AI model and flows through server-side execution to the client-side UI.
Sources: lib/ai/tools/create-document.ts:16-30, components/ai-elements/tool.tsx:28-44
Server-Side Implementation
To add a new tool, you define it using the tool function from the ai package. Every tool requires a description, an inputSchema, and an execute function.
| Parameter | Type | Description |
|---|---|---|
description | string | Explains to the AI model when and how to use the tool. |
inputSchema | ZodObject | Defines the parameters the AI must provide. |
execute | Function | The asynchronous logic performed by the tool. |
Sources: lib/ai/tools/create-document.ts:20-30, lib/ai/tools/update-document.ts:18-24
Data Streaming
Tools use the UIMessageStreamWriter to send transient data to the UI before the tool execution completes. This allows the UI to update in real-time, such as clearing previous content or setting document titles.
dataStream.write({
type: "data-title",
data: title,
transient: true,
});
Sources: lib/ai/tools/create-document.ts:46-50, lib/ai/tools/update-document.ts:37-41
Client-Side Rendering
The system uses standard UI components to display tool calls and results. Tools transition through several states during their lifecycle, which are reflected in the UI via the ToolHeader and ToolOutput components.
Tool Lifecycle States
| State | Label | Description |
|---|---|---|
input-streaming | Pending | Tool arguments are being generated by the model. |
input-available | Running | Tool execution has started on the server. |
output-available | Completed | Tool execution finished successfully. |
output-error | Error | An error occurred during execution. |
approval-requested | Awaiting Approval | The tool requires manual user confirmation. |
Sources: components/ai-elements/tool.tsx:46-54
Rendering Tool Parts
The PreviewMessage component in components/chat/message.tsx handles the switching logic for different tool types. For example, tool-createDocument renders a DocumentPreview, while tool-requestSuggestions uses the generic Tool wrapper.
Tool Integration Steps
- Define Tool Logic: Create a file in
lib/ai/tools/and export the tool definition. - Register Types: Add the tool to the
ChatToolstype inlib/types.ts. - Update Prompt: Include the tool's description and usage instructions in
lib/ai/prompts.tsso the model knows when to use it. - Handle UI Rendering: Add a conditional block in
components/chat/message.tsxto handle thetool-yourToolNamepart type.
Sources: lib/types.ts:24-31, lib/ai/prompts.ts:47-65, components/chat/message.tsx:247-280
Example: Document Suggestions Tool
The requestSuggestions tool demonstrates how to handle complex outputs. It streams individual suggestions to the client using the data-suggestion data type and saves them to the database.
- Input:
documentId(UUID) - Execution: Reads document content, calls a secondary AI stream to generate improvements, and writes suggestions to the data stream.
- Output: Returns a confirmation message and the document ID.
Sources: lib/ai/tools/request-suggestions.ts:20-109
Conclusion
Adding custom AI tools involves defining server-side execution logic with Zod schemas and mapping those calls to React components. By utilizing the UIMessageStreamWriter, tools provide a responsive experience through real-time updates and clear status indicators in the chat interface.