AI SDK Persistence DB

The application uses a prefix-based approach to persist chat messages in PostgreSQL, avoiding complex polymorphic relationships while maintaining type safety and performance.

Schema Design

The system relies on three primary tables:

• chats: Stores chat sessions with auto-generated IDs.
• messages: Stores individual messages, including their role, timestamp, and a reference to the parent chat.
• parts: Stores the content of message parts using a prefix-based column naming convention.

Prefix Convention

Each message part type uses dedicated columns with specific prefixes to ensure structured storage:

• text_*: Text content.
• reasoning_*: Reasoning or thinking process.
• file_*: File attachments.
• source_url_*: URL sources.
• source_document_*: Document sources.
• tool_[toolName]_*: Tool calls (e.g., tool_getWeatherInformation_*).
• data_[dataType]_*: Custom data parts (e.g., data_weather_*).

Data Integrity

The schema enforces integrity through constraints that require complete part definitions. For example, a source URL part must include both source_url_url and source_url_description. Tool calls must also provide a complete set of columns (input, output, state, etc.) to prevent partial or corrupted data.

The persistence layer uses a PostgreSQL schema managed via Drizzle ORM, organized into three primary tables: chats, messages, and parts. This structure allows for hierarchical storage of conversations, where a chat contains multiple messages, and each message contains multiple ordered parts (text, reasoning, files, tool calls, etc.).

Table Overview

chats

Represents a single conversation session.

• id: Primary key (generated via generateId()).

messages

Represents an individual message within a chat.

• id: Primary key.
• chatId: Foreign key referencing chats.id (with onDelete: 'cascade').
• role: The role of the message (e.g., user, assistant, system).
• createdAt: Timestamp of message creation.

parts

Represents the granular content components of a message. This table uses a single-table inheritance pattern where different type values determine which columns are populated.

• id: Primary key.
• messageId: Foreign key referencing messages.id (with onDelete: 'cascade').
• type: The type of part (e.g., text, reasoning, file, tool-getLocation, data-weather).
• order: Integer used to maintain the sequence of parts within a message.
• Content Columns: Specific columns are used based on the type (e.g., text_text for text parts, file_url for file parts, or tool_state for tool parts).
• Constraints: The schema enforces data integrity via SQL check constraints, ensuring that if a part is of a specific type, its required fields are not null.

When using the AI SDK with persistence, you can define custom types for your messages to include specific metadata, data parts (for structured UI state), and tool sets. This ensures type safety when retrieving and storing messages in your database.

To define your own message types, you should compose UIMessage and UIMessagePart using your custom Zod schemas for metadata and data parts, and the InferUITools utility for your tool definitions.

import { InferUITools, JSONValue, UIMessage, UIMessagePart } from "ai";
import z from "zod";

// 1. Define your metadata schema
export const metadataSchema = z.object({});
type MyMetadata = z.infer<typeof metadataSchema>;

// 2. Define your data part schema (e.g., for weather widgets or other UI state)
export const dataPartSchema = z.object({
  weather: z.object({
    weather: z.string().optional(),
    location: z.string().optional(),
    temperature: z.number().optional(),
    loading: z.boolean().default(true),
  }),
});
type MyDataPart = z.infer<typeof dataPartSchema>;

// 3. Define your tool set (assuming 'tools' is imported from your tool definitions)
export type MyToolSet = InferUITools<ReturnType<typeof tools>>;

// 4. Compose the final UIMessage and UIMessagePart types
export type MyUIMessage = UIMessage<MyMetadata, MyDataPart, MyToolSet>;
export type MyUIMessagePart = UIMessagePart<MyDataPart, MyToolSet>;

The parts table uses specific columns to store different types of AI SDK message parts. The type column dictates which of these columns must be populated according to database check constraints.

// Text & Reasoning
text_text: text(),
reasoning_text: text(),

// File
file_mediaType: varchar(),
file_filename: varchar(),
file_url: varchar(),

// Source URL
source_url_sourceId: varchar(),
source_url_url: varchar(),
source_url_title: varchar(),

// Source Document
source_document_sourceId: varchar(),
source_document_mediaType: varchar(),
source_document_title: varchar(),
source_document_filename: varchar(),

// Tool Calls
tool_toolCallId: varchar(),
tool_state: varchar(),
tool_errorText: varchar(),
tool_getWeatherInformation_input: jsonb(),
tool_getWeatherInformation_output: jsonb(),
tool_getLocation_input: jsonb(),
tool_getLocation_output: jsonb(),

// Data Parts
data_weather_id: varchar(),
data_weather_location: varchar(),
data_weather_weather: varchar(),
data_weather_temperature: real(),

// Metadata
providerMetadata: jsonb(),